From ceb51dda4a6c48fced3e9d3df333b1702b1f83db Mon Sep 17 00:00:00 2001 From: Thiago Crepaldi Date: Thu, 30 Sep 2021 10:37:35 -0400 Subject: [PATCH] Support external torch cpp extensions on ORTModule (#9223) --- cmake/onnxruntime_python.cmake | 19 ++-- .../python/training/ortmodule/__init__.py | 16 ++-- .../_custom_autograd_function_runner.py | 2 +- .../ortmodule/_graph_execution_manager.py | 5 +- .../torch_cpp_extensions/__init__.py | 57 ++++------- .../torch_cpp_extensions/cpu/__init__.py | 0 .../cpu/aten_op_executor/__init__.py | 37 ++++++++ .../aten_op_executor/aten_op_executor.cc | 0 .../{ => cpu}/aten_op_executor/setup.py | 0 .../{ => cpu}/torch_interop_utils/setup.py | 0 .../torch_interop_utils.cc | 0 .../torch_cpp_extensions/cuda/__init__.py | 0 .../{ => cuda}/torch_gpu_allocator/setup.py | 10 +- .../torch_gpu_allocator.cc | 0 .../ortmodule/torch_cpp_extensions/install.py | 94 ++++++++++--------- .../orttraining_test_ortmodule_fallback.py | 4 +- setup.py | 12 +-- 17 files changed, 138 insertions(+), 118 deletions(-) create mode 100644 orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/__init__.py create mode 100644 orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/__init__.py rename orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/{ => cpu}/aten_op_executor/aten_op_executor.cc (100%) rename orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/{ => cpu}/aten_op_executor/setup.py (100%) rename orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/{ => cpu}/torch_interop_utils/setup.py (100%) rename orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/{ => cpu}/torch_interop_utils/torch_interop_utils.cc (100%) create mode 100644 orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/__init__.py rename orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/{ => cuda}/torch_gpu_allocator/setup.py (83%) rename orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/{ => cuda}/torch_gpu_allocator/torch_gpu_allocator.cc (100%) diff --git a/cmake/onnxruntime_python.cmake b/cmake/onnxruntime_python.cmake index cac8b42f9e..b75111fee9 100644 --- a/cmake/onnxruntime_python.cmake +++ b/cmake/onnxruntime_python.cmake @@ -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 $/onnxruntime/training/ortmodule/experimental/json_config COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/training/ortmodule/experimental/hierarchical_ortmodule COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/training/ortmodule/torch_cpp_extensions - COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/training/ortmodule/torch_cpp_extensions/aten_op_executor - COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/training/ortmodule/torch_cpp_extensions/torch_interop_utils - COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/training/ortmodule/torch_cpp_extensions/torch_gpu_allocator + COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor + COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils + COMMAND ${CMAKE_COMMAND} -E make_directory $/onnxruntime/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator COMMAND ${CMAKE_COMMAND} -E copy ${onnxruntime_python_capi_training_srcs} $/onnxruntime/capi/training/ @@ -554,13 +554,13 @@ if (onnxruntime_ENABLE_TRAINING) $/onnxruntime/training/ortmodule/torch_cpp_extensions/ COMMAND ${CMAKE_COMMAND} -E copy ${onnxruntime_python_ortmodule_torch_cpp_ext_aten_op_executor_srcs} - $/onnxruntime/training/ortmodule/torch_cpp_extensions/aten_op_executor/ + $/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} - $/onnxruntime/training/ortmodule/torch_cpp_extensions/torch_interop_utils/ + $/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} - $/onnxruntime/training/ortmodule/torch_cpp_extensions/torch_gpu_allocator/ + $/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() - diff --git a/orttraining/orttraining/python/training/ortmodule/__init__.py b/orttraining/orttraining/python/training/ortmodule/__init__.py index 660a4b26e5..0fd1460e2f 100644 --- a/orttraining/orttraining/python/training/ortmodule/__init__.py +++ b/orttraining/orttraining/python/training/ortmodule/__init__.py @@ -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 diff --git a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_runner.py b/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_runner.py index 67bef86af1..ae32b0304c 100644 --- a/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_runner.py +++ b/orttraining/orttraining/python/training/ortmodule/_custom_autograd_function_runner.py @@ -105,7 +105,7 @@ def call_python_forward_function( # output_1, output_2 --- shared_ptr --- ↑ # ↓ 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). diff --git a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py index 8de45c29e5..8d5af41297 100644 --- a/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py +++ b/orttraining/orttraining/python/training/ortmodule/_graph_execution_manager.py @@ -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, diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/__init__.py b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/__init__.py index df7ad283f7..dbb5e4fcee 100644 --- a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/__init__.py +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/__init__.py @@ -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 diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/__init__.py b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/__init__.py b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/__init__.py new file mode 100644 index 0000000000..3b58c55051 --- /dev/null +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/__init__.py @@ -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 diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/aten_op_executor/aten_op_executor.cc b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/aten_op_executor.cc similarity index 100% rename from orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/aten_op_executor/aten_op_executor.cc rename to orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/aten_op_executor.cc diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/aten_op_executor/setup.py b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/setup.py similarity index 100% rename from orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/aten_op_executor/setup.py rename to orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/aten_op_executor/setup.py diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/torch_interop_utils/setup.py b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils/setup.py similarity index 100% rename from orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/torch_interop_utils/setup.py rename to orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils/setup.py diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/torch_interop_utils/torch_interop_utils.cc b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils/torch_interop_utils.cc similarity index 100% rename from orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/torch_interop_utils/torch_interop_utils.cc rename to orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cpu/torch_interop_utils/torch_interop_utils.cc diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/__init__.py b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/torch_gpu_allocator/setup.py b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator/setup.py similarity index 83% rename from orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/torch_gpu_allocator/setup.py rename to orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator/setup.py index 0d556c6af5..6a36e16699 100644 --- a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/torch_gpu_allocator/setup.py +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator/setup.py @@ -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__), diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/torch_gpu_allocator/torch_gpu_allocator.cc b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator/torch_gpu_allocator.cc similarity index 100% rename from orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/torch_gpu_allocator/torch_gpu_allocator.cc rename to orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/cuda/torch_gpu_allocator/torch_gpu_allocator.cc diff --git a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/install.py b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/install.py index 3818059d9e..21282bcce2 100644 --- a/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/install.py +++ b/orttraining/orttraining/python/training/ortmodule/torch_cpp_extensions/install.py @@ -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() diff --git a/orttraining/orttraining/test/python/orttraining_test_ortmodule_fallback.py b/orttraining/orttraining/test/python/orttraining_test_ortmodule_fallback.py index d4cab55f74..e130794980 100644 --- a/orttraining/orttraining/test/python/orttraining_test_ortmodule_fallback.py +++ b/orttraining/orttraining/test/python/orttraining_test_ortmodule_fallback.py @@ -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: diff --git a/setup.py b/setup.py index bd1dec2699..145064dde7 100644 --- a/setup.py +++ b/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: