From 1e13e2666e61afcdc82ff17253555f5a8da65f86 Mon Sep 17 00:00:00 2001 From: Weixing Zhang Date: Wed, 10 Mar 2021 10:00:35 -0800 Subject: [PATCH] Support ROCM EP for ORTModule (#6967) 1. Disable external allocator for ROCM EP since it is not supported yet. 2. For AMD GPU, the EP name is ROCMExecutionProvider --- orttraining/orttraining/python/training/ortmodule.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/orttraining/orttraining/python/training/ortmodule.py b/orttraining/orttraining/python/training/ortmodule.py index 32dd024916..1487757a11 100644 --- a/orttraining/orttraining/python/training/ortmodule.py +++ b/orttraining/orttraining/python/training/ortmodule.py @@ -256,8 +256,12 @@ class ORTModule(torch.nn.Module): self._save_onnx = False self._save_onnx_prefix = '' + from torch.utils.cpp_extension import ROCM_HOME + self.is_rocm_pytorch = (True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False) + # CPP extension to get torch CUDA allocator's alloc and free function addresses - self._use_external_cuda_allocator = True + # Disable external allocator for ROCM EP since external allocator is not supported yet. + self._use_external_cuda_allocator = (False if self.is_rocm_pytorch else True) if self._use_external_cuda_allocator: self._torch_cuda_allocator = _load_torch_allocator_cpp_extension() self._torch_alloc = self._torch_cuda_allocator.cuda_caching_allocator_raw_alloc_address() @@ -289,7 +293,8 @@ class ORTModule(torch.nn.Module): provider_options = None if self._device.type == 'cuda': # Configure the InferenceSessions to use the specific GPU on which the model is placed. - providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] + providers = (["ROCMExecutionProvider"] if self.is_rocm_pytorch else ["CUDAExecutionProvider"]) + providers.append("CPUExecutionProvider") if self._use_external_cuda_allocator: provider_options = [{"device_id": str(self._device.index), "cuda_external_alloc": str(self._torch_alloc), "cuda_external_free": str(self._torch_free)}, {}] else: