mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-19 19:00:47 +00:00
Support external torch cpp extensions on ORTModule (#9223)
This commit is contained in:
Thiago Crepaldi
GitHub
parent
ffca0b777b
commit
ceb51dda4a
17 changed files with 138 additions and 118 deletions
|
|
@ -281,13 +281,13 @@ if (onnxruntime_ENABLE_TRAINING)
|
|||
"${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/torch_cpp_extensions/*.py"
|
||||
)
|
||||
file(GLOB onnxruntime_python_ortmodule_torch_cpp_ext_aten_op_executor_srcs CONFIGURE_DEPENDS
|
||||
"${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/torch_cpp_extensions/aten_op_executor/*"
|
||||
"${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/*"
|
||||
)
|
||||
file(GLOB onnxruntime_python_ortmodule_torch_cpp_ext_torch_interop_utils_srcs CONFIGURE_DEPENDS
|
||||
"${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/torch_cpp_extensions/torch_interop_utils/*"
|
||||
"${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils/*"
|
||||
)
|
||||
file(GLOB onnxruntime_python_ortmodule_torch_cpp_ext_torch_gpu_allocator_srcs CONFIGURE_DEPENDS
|
||||
"${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/torch_cpp_extensions/torch_gpu_allocator/*"
|
||||
"${ORTTRAINING_SOURCE_DIR}/python/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator/*"
|
||||
)
|
||||
else()
|
||||
file(GLOB onnxruntime_python_capi_training_srcs CONFIGURE_DEPENDS
|
||||
|
|
@ -522,9 +522,9 @@ if (onnxruntime_ENABLE_TRAINING)
|
|||
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/experimental/json_config
|
||||
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/experimental/hierarchical_ortmodule
|
||||
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions
|
||||
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/aten_op_executor
|
||||
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/torch_interop_utils
|
||||
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/torch_gpu_allocator
|
||||
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor
|
||||
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils
|
||||
COMMAND ${CMAKE_COMMAND} -E make_directory $<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator
|
||||
COMMAND ${CMAKE_COMMAND} -E copy
|
||||
${onnxruntime_python_capi_training_srcs}
|
||||
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/capi/training/
|
||||
|
|
@ -554,13 +554,13 @@ if (onnxruntime_ENABLE_TRAINING)
|
|||
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/
|
||||
COMMAND ${CMAKE_COMMAND} -E copy
|
||||
${onnxruntime_python_ortmodule_torch_cpp_ext_aten_op_executor_srcs}
|
||||
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/aten_op_executor/
|
||||
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/
|
||||
COMMAND ${CMAKE_COMMAND} -E copy
|
||||
${onnxruntime_python_ortmodule_torch_cpp_ext_torch_interop_utils_srcs}
|
||||
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/torch_interop_utils/
|
||||
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils/
|
||||
COMMAND ${CMAKE_COMMAND} -E copy
|
||||
${onnxruntime_python_ortmodule_torch_cpp_ext_torch_gpu_allocator_srcs}
|
||||
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/torch_gpu_allocator/
|
||||
$<TARGET_FILE_DIR:${build_output_target}>/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator/
|
||||
)
|
||||
endif()
|
||||
|
||||
|
|
@ -657,4 +657,3 @@ endif()
|
|||
if (onnxruntime_ENABLE_LANGUAGE_INTEROP_OPS)
|
||||
include(onnxruntime_language_interop_ops.cmake)
|
||||
endif()
|
||||
|
||||
|
|
|
|||
|
|
@ -6,23 +6,22 @@
|
|||
import os
|
||||
import sys
|
||||
import torch
|
||||
from packaging import version
|
||||
|
||||
from onnxruntime import set_seed
|
||||
from packaging import version
|
||||
from onnxruntime.capi import build_and_package_info as ort_info
|
||||
from ._fallback import (_FallbackPolicy,
|
||||
ORTModuleFallbackException,
|
||||
ORTModuleInitException,
|
||||
wrap_exception)
|
||||
from .torch_cpp_extensions import is_installed as is_torch_cpp_extensions_installed
|
||||
|
||||
|
||||
################################################################################
|
||||
# All global constant goes here, before ORTModule is imported ##################
|
||||
################################################################################
|
||||
ONNX_OPSET_VERSION = 12
|
||||
MINIMUM_RUNTIME_PYTORCH_VERSION_STR = '1.8.1'
|
||||
TORCH_CPP_DIR = os.path.join(os.path.dirname(__file__),
|
||||
'torch_cpp_extensions')
|
||||
ORTMODULE_TORCH_CPP_DIR = os.path.join(os.path.dirname(__file__), 'torch_cpp_extensions')
|
||||
_FALLBACK_INIT_EXCEPTION = None
|
||||
ORTMODULE_FALLBACK_POLICY = _FallbackPolicy.FALLBACK_UNSUPPORTED_DEVICE |\
|
||||
_FallbackPolicy.FALLBACK_UNSUPPORTED_DATA |\
|
||||
|
|
@ -32,6 +31,9 @@ ORTMODULE_FALLBACK_POLICY = _FallbackPolicy.FALLBACK_UNSUPPORTED_DEVICE |\
|
|||
ORTMODULE_FALLBACK_RETRY = False
|
||||
ORTMODULE_IS_DETERMINISTIC = torch.are_deterministic_algorithms_enabled()
|
||||
|
||||
ONNXRUNTIME_CUDA_VERSION = ort_info.cuda_version if hasattr(ort_info, 'cuda_version') else ''
|
||||
ONNXRUNTIME_ROCM_VERSION = ort_info.rocm_version if hasattr(ort_info, 'rocm_version') else ''
|
||||
|
||||
# Verify minimum PyTorch version is installed before proceding to ONNX Runtime initialization
|
||||
try:
|
||||
import torch
|
||||
|
|
@ -51,12 +53,12 @@ except ImportError as e:
|
|||
'installed in order to run ONNX Runtime ORTModule frontend!') from e
|
||||
|
||||
# Verify whether PyTorch C++ extensions are already compiled
|
||||
|
||||
if not is_torch_cpp_extensions_installed(TORCH_CPP_DIR) and '-m' not in sys.argv:
|
||||
# TODO: detect when installed extensions are outdated and need reinstallation. Hash? Version file?
|
||||
if not is_torch_cpp_extensions_installed(ORTMODULE_TORCH_CPP_DIR) and '-m' not in sys.argv:
|
||||
_FALLBACK_INIT_EXCEPTION = wrap_exception(
|
||||
ORTModuleInitException,
|
||||
EnvironmentError(
|
||||
f"ORTModule's extensions were not detected at '{TORCH_CPP_DIR}' folder. "
|
||||
f"ORTModule's extensions were not detected at '{ORTMODULE_TORCH_CPP_DIR}' folder. "
|
||||
"Run `python -m torch_ort.configure` before using `ORTModule` frontend."))
|
||||
|
||||
# Initalized ORT's random seed with pytorch's initial seed
|
||||
|
|
|
|||
|
|
@ -105,7 +105,7 @@ def call_python_forward_function(
|
|||
# output_1, output_2 --- shared_ptr<PyNode> --- ↑
|
||||
# ↓ previous gradient function
|
||||
|
||||
# We remove the edges starting between current autograd.Function's gradient function and
|
||||
# We remove the edges starting between current autograd.Function's gradient function and
|
||||
# it's input's gradient function (e.g. AccumulateGrad gradient function), then
|
||||
# AccumulateGrad gradient function will be destroyed, releasing the reference to input_1
|
||||
# (https://github.com/pytorch/pytorch/blob/15532595209d2daf34d35e10f8d3d3b64966aea2/torch/csrc/autograd/functions/accumulate_grad.cpp#L21).
|
||||
|
|
|
|||
|
|
@ -7,9 +7,9 @@ from .debug_options import DebugOptions, LogLevel
|
|||
from . import (_utils,
|
||||
_io,
|
||||
_logger,
|
||||
torch_cpp_extensions as _cpp_ext,
|
||||
_onnx_models,
|
||||
_are_deterministic_algorithms_enabled)
|
||||
from .torch_cpp_extensions.cpu.aten_op_executor import load_aten_op_executor_cpp_extension_if_needed
|
||||
from ._custom_autograd_function import custom_autograd_function_enabler
|
||||
from ._custom_autograd_function_exporter import _post_process_after_export
|
||||
from ._graph_execution_interface import GraphExecutionInterface
|
||||
|
|
@ -302,8 +302,7 @@ class GraphExecutionManager(GraphExecutionInterface):
|
|||
self._set_device_from_module(inputs, kwargs)
|
||||
self._onnx_models.exported_model = self._get_exported_model(
|
||||
schema, *inputs, **kwargs)
|
||||
_cpp_ext._load_aten_op_executor_cpp_extension_if_needed(
|
||||
self._onnx_models.exported_model)
|
||||
load_aten_op_executor_cpp_extension_if_needed(self._onnx_models.exported_model)
|
||||
if self._debug_options.save_onnx_models.save:
|
||||
self._onnx_models.save_exported_model(self._debug_options.save_onnx_models.path,
|
||||
self._debug_options.save_onnx_models.name_prefix,
|
||||
|
|
|
|||
|
|
@ -5,16 +5,31 @@
|
|||
|
||||
"""Support for PyTorch C++ extensions within ORTModule
|
||||
|
||||
Pytorch C++ extensions must be added to this (ORTMODULE_TORCH_CPP_DIR) folder to be automatically
|
||||
detected and installed by `python -m torch_ort.configure`
|
||||
|
||||
Each extension must be within a folder and contain a setup.py file
|
||||
|
||||
CUDA extensions must be stored within 'cuda' folders
|
||||
CPU extensions must be stored within 'cpu' subfolder
|
||||
|
||||
Extensions are lexicographically ordered for compilation.
|
||||
e.g. '001_my_extension' is compiled before '002_my_other_extension'
|
||||
|
||||
The following environment variables are available for the extensions setup.py
|
||||
|
||||
- ORTMODULE_TORCH_CPP_DIR: ORTModule's internal
|
||||
- ONNXRUNTIME_ROCM_VERSION: ROCM version used to build ONNX Runtime package
|
||||
- ONNXRUNTIME_CUDA_VERSION: CUDA version used to build ONNX Runtime package
|
||||
|
||||
TODO: Create a generic mechanism to pass arguments from ORTModule into each extension setup.py
|
||||
TODO: Create environment variables to allow extensions to be hosted outside ONNX runtime installation folder
|
||||
(e.g. ORTMODULE_EXTERNAL_TORCH_CPP_EXTENSION_DIR, ORTMODULE_EXTERNAL_TORCH_CUDA_EXTENSION_DIR)
|
||||
|
||||
TODO: Implement mechanism to register extensions and prevent issues with incorrect/missing flags
|
||||
for each :meth:`torch.utils.cpp_extension.*` call
|
||||
"""
|
||||
|
||||
import os
|
||||
import threading
|
||||
from functools import wraps
|
||||
from glob import glob
|
||||
from onnxruntime.capi import _pybind_state as C
|
||||
|
||||
|
||||
def is_installed(torch_cpp_extension_path):
|
||||
|
|
@ -22,35 +37,3 @@ def is_installed(torch_cpp_extension_path):
|
|||
torch_cpp_exts.extend(glob(os.path.join(torch_cpp_extension_path, '*.dll')))
|
||||
torch_cpp_exts.extend(glob(os.path.join(torch_cpp_extension_path, '*.dylib')))
|
||||
return len(torch_cpp_exts) > 0
|
||||
|
||||
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():
|
||||
from onnxruntime.training.ortmodule.torch_cpp_extensions import aten_op_executor
|
||||
C.register_aten_op_executor(str(aten_op_executor.is_tensor_argument_address()),
|
||||
str(aten_op_executor.execute_aten_operator_address()))
|
||||
|
||||
|
||||
def _load_aten_op_executor_cpp_extension_if_needed(onnx_model):
|
||||
for node in onnx_model.graph.node:
|
||||
if node.op_type == 'ATenOp' and node.domain == 'com.microsoft':
|
||||
_load_aten_op_executor_cpp_extension()
|
||||
break
|
||||
|
|
|
|||
|
|
@ -0,0 +1,37 @@
|
|||
from onnxruntime.capi import _pybind_state as C
|
||||
|
||||
import threading
|
||||
from functools import wraps
|
||||
|
||||
|
||||
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():
|
||||
from onnxruntime.training.ortmodule.torch_cpp_extensions import aten_op_executor
|
||||
C.register_aten_op_executor(str(aten_op_executor.is_tensor_argument_address()),
|
||||
str(aten_op_executor.execute_aten_operator_address()))
|
||||
|
||||
|
||||
def load_aten_op_executor_cpp_extension_if_needed(onnx_model):
|
||||
for node in onnx_model.graph.node:
|
||||
if node.op_type == 'ATenOp' and node.domain == 'com.microsoft':
|
||||
_load_aten_op_executor_cpp_extension()
|
||||
break
|
||||
|
|
@ -11,16 +11,8 @@ from setuptools import setup
|
|||
from torch.utils import cpp_extension
|
||||
|
||||
|
||||
def parse_arg_remove_boolean(argv, arg_name):
|
||||
arg_value = False
|
||||
if arg_name in sys.argv:
|
||||
arg_value = True
|
||||
argv.remove(arg_name)
|
||||
|
||||
return arg_value
|
||||
|
||||
# TODO: Implement a cleaner way to auto-generate torch_gpu_allocator.cc
|
||||
use_rocm = True if parse_arg_remove_boolean(sys.argv, '--use_rocm') else False
|
||||
use_rocm = True if os.environ['ONNXRUNTIME_ROCM_VERSION'] else False
|
||||
gpu_identifier = "hip" if use_rocm else "cuda"
|
||||
gpu_allocator_header = "HIPCachingAllocator" if use_rocm else "CUDACachingAllocator"
|
||||
filename = os.path.join(os.path.dirname(__file__),
|
||||
|
|
@ -3,85 +3,93 @@
|
|||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
from onnxruntime.capi import build_and_package_info as ort_info
|
||||
from onnxruntime.training.ortmodule import (ORTMODULE_TORCH_CPP_DIR,
|
||||
ONNXRUNTIME_CUDA_VERSION,
|
||||
ONNXRUNTIME_ROCM_VERSION)
|
||||
|
||||
from glob import glob
|
||||
from shutil import copyfile
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
|
||||
from glob import glob
|
||||
from shutil import copyfile
|
||||
|
||||
def _list_extensions(path):
|
||||
extensions = []
|
||||
for root, _, files in os.walk(path):
|
||||
for name in files:
|
||||
if name.lower() == 'setup.py':
|
||||
extensions.append(os.path.join(root, name))
|
||||
return extensions
|
||||
|
||||
|
||||
def _list_cpu_extensions():
|
||||
return _list_extensions(os.path.join(ORTMODULE_TORCH_CPP_DIR, 'cpu'))
|
||||
|
||||
|
||||
def _list_cuda_extensions():
|
||||
return _list_extensions(os.path.join(ORTMODULE_TORCH_CPP_DIR, 'cuda'))
|
||||
|
||||
|
||||
def _install_extension(ext_name, ext_path, cwd):
|
||||
ret_code = subprocess.call(f"{sys.executable} {ext_path} build",
|
||||
cwd=cwd,
|
||||
shell=True)
|
||||
if ret_code != 0:
|
||||
print(f'There was an error compiling "{ext_name}" PyTorch CPP extension')
|
||||
sys.exit(ret_code)
|
||||
|
||||
|
||||
def build_torch_cpp_extensions():
|
||||
'''Builds PyTorch CPP extensions and returns metadata'''
|
||||
|
||||
cuda_version = ort_info.cuda_version if hasattr(ort_info, 'cuda_version') else None
|
||||
rocm_version = ort_info.rocm_version if hasattr(ort_info, 'rocm_version') else None
|
||||
|
||||
# Run this from within onnxruntime package folder
|
||||
is_gpu_available = cuda_version is not None or rocm_version is not None
|
||||
cpp_ext_dir = os.path.join(os.path.dirname(__file__))
|
||||
os.chdir(cpp_ext_dir)
|
||||
is_gpu_available = ONNXRUNTIME_CUDA_VERSION is not None or ONNXRUNTIME_ROCM_VERSION is not None
|
||||
os.chdir(ORTMODULE_TORCH_CPP_DIR)
|
||||
|
||||
# Extensions might leverage CUDA/ROCM versions internally
|
||||
os.environ["ONNXRUNTIME_CUDA_VERSION"] = ONNXRUNTIME_CUDA_VERSION
|
||||
os.environ["ONNXRUNTIME_ROCM_VERSION"] = ONNXRUNTIME_ROCM_VERSION
|
||||
|
||||
############################################################################
|
||||
# Pytorch CPP Extensions that DO require CUDA/ROCM
|
||||
############################################################################
|
||||
if is_gpu_available:
|
||||
setup_script = os.path.join(cpp_ext_dir,
|
||||
'torch_gpu_allocator',
|
||||
'setup.py')
|
||||
version = '--use_rocm' if rocm_version else ''
|
||||
ret_code = subprocess.call(f"{sys.executable} {setup_script} build {version}",
|
||||
cwd=cpp_ext_dir,
|
||||
shell=True)
|
||||
if ret_code != 0:
|
||||
print('There was an error compiling "torch_gpu_allocator" PyTorch CPP extension')
|
||||
sys.exit(ret_code)
|
||||
for ext_setup in _list_cuda_extensions():
|
||||
_install_extension(ext_setup.split(
|
||||
os.sep)[-2], ext_setup, ORTMODULE_TORCH_CPP_DIR)
|
||||
|
||||
############################################################################
|
||||
# Pytorch CPP Extensions that DO NOT require CUDA/ROCM
|
||||
############################################################################
|
||||
setup_script = os.path.join(cpp_ext_dir,
|
||||
'aten_op_executor',
|
||||
'setup.py')
|
||||
ret_code = subprocess.call(f"{sys.executable} {setup_script} build",
|
||||
cwd=cpp_ext_dir,
|
||||
shell=True)
|
||||
if ret_code != 0:
|
||||
print('There was an error compiling "aten_op_executor" PyTorch CPP extension')
|
||||
sys.exit(ret_code)
|
||||
|
||||
setup_script = os.path.join(cpp_ext_dir,
|
||||
'torch_interop_utils',
|
||||
'setup.py')
|
||||
ret_code = subprocess.call(f"{sys.executable} {setup_script} build",
|
||||
cwd=cpp_ext_dir,
|
||||
shell=True)
|
||||
if ret_code != 0:
|
||||
print('There was an error compiling "torch_interop_utils" PyTorch CPP extension')
|
||||
sys.exit(ret_code)
|
||||
for ext_setup in _list_cpu_extensions():
|
||||
_install_extension(ext_setup.split(
|
||||
os.sep)[-2], ext_setup, ORTMODULE_TORCH_CPP_DIR)
|
||||
|
||||
############################################################################
|
||||
# Copy Pytorch CPP Extensions to the local onnxruntime package folder
|
||||
# Install Pytorch CPP Extensions into local onnxruntime package folder
|
||||
############################################################################
|
||||
torch_cpp_exts = glob(os.path.join(cpp_ext_dir,
|
||||
torch_cpp_exts = glob(os.path.join(ORTMODULE_TORCH_CPP_DIR,
|
||||
'build',
|
||||
'lib.*',
|
||||
'*.so'))
|
||||
torch_cpp_exts.extend(glob(os.path.join(cpp_ext_dir,
|
||||
torch_cpp_exts.extend(glob(os.path.join(ORTMODULE_TORCH_CPP_DIR,
|
||||
'build',
|
||||
'lib.*',
|
||||
'*.dll')))
|
||||
torch_cpp_exts.extend(glob(os.path.join(cpp_ext_dir,
|
||||
torch_cpp_exts.extend(glob(os.path.join(ORTMODULE_TORCH_CPP_DIR,
|
||||
'build',
|
||||
'lib.*',
|
||||
'*.dylib')))
|
||||
for ext in torch_cpp_exts:
|
||||
dest_ext = os.path.join(cpp_ext_dir, os.path.basename(ext))
|
||||
dest_ext = os.path.join(ORTMODULE_TORCH_CPP_DIR, os.path.basename(ext))
|
||||
print(f'Installing {ext} -> {dest_ext}')
|
||||
copyfile(ext, dest_ext)
|
||||
|
||||
# Tear down
|
||||
os.environ.pop("ONNXRUNTIME_CUDA_VERSION")
|
||||
os.environ.pop("ONNXRUNTIME_ROCM_VERSION")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
build_torch_cpp_extensions()
|
||||
|
|
|
|||
|
|
@ -9,7 +9,7 @@ import torch
|
|||
import pytest
|
||||
import warnings
|
||||
|
||||
from onnxruntime.training.ortmodule import ORTModule, _fallback, TORCH_CPP_DIR
|
||||
from onnxruntime.training.ortmodule import ORTModule, _fallback, ORTMODULE_TORCH_CPP_DIR
|
||||
from onnxruntime.training.ortmodule.torch_cpp_extensions import is_installed as is_torch_cpp_extensions_installed
|
||||
import _test_helpers
|
||||
from _orttraining_ortmodule_models import (NeuralNetSinglePositionalArgument,
|
||||
|
|
@ -387,7 +387,7 @@ def test_ortmodule_fallback_init__missing_cpp_extensions(is_training, fallback_e
|
|||
# matching_policy: True matches FALLBACK_UNSUPPORTED_TORCH_MODEL policy to ORTModuleDeviceException exception.
|
||||
# Otherwise, an incorrect policy (FALLBACK_UNSUPPORTED_DEVICE) is used to verify that the fallback does not happen
|
||||
|
||||
if is_torch_cpp_extensions_installed(TORCH_CPP_DIR):
|
||||
if is_torch_cpp_extensions_installed(ORTMODULE_TORCH_CPP_DIR):
|
||||
warnings.warn('Skipping test_ortmodule_fallback_init__missing_cpp_extensions.'
|
||||
f' It requires PyTorch CPP extensions to be missing')
|
||||
else:
|
||||
|
|
|
|||
12
setup.py
12
setup.py
|
|
@ -329,12 +329,12 @@ if enable_training:
|
|||
'onnxruntime.training.ortmodule.experimental.json_config',
|
||||
'onnxruntime.training.ortmodule.experimental.hierarchical_ortmodule',
|
||||
'onnxruntime.training.ortmodule.torch_cpp_extensions',
|
||||
'onnxruntime.training.ortmodule.torch_cpp_extensions.aten_op_executor',
|
||||
'onnxruntime.training.ortmodule.torch_cpp_extensions.torch_interop_utils',
|
||||
'onnxruntime.training.ortmodule.torch_cpp_extensions.torch_gpu_allocator'])
|
||||
package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.aten_op_executor'] = ['*.cc']
|
||||
package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.torch_interop_utils'] = ['*.cc']
|
||||
package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.torch_gpu_allocator'] = ['*.cc']
|
||||
'onnxruntime.training.ortmodule.torch_cpp_extensions.cpu.aten_op_executor',
|
||||
'onnxruntime.training.ortmodule.torch_cpp_extensions.cpu.torch_interop_utils',
|
||||
'onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.torch_gpu_allocator'])
|
||||
package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.cpu.aten_op_executor'] = ['*.cc']
|
||||
package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.cpu.torch_interop_utils'] = ['*.cc']
|
||||
package_data['onnxruntime.training.ortmodule.torch_cpp_extensions.cuda.torch_gpu_allocator'] = ['*.cc']
|
||||
requirements_file = "requirements-training.txt"
|
||||
# with training, we want to follow this naming convention:
|
||||
# stable:
|
||||
|
|
|
|||
Loading…
Reference in a new issue