saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-30 20:18:08 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Support ROCM EP for ORTModule (#6967)

Browse source 
1. Disable external allocator for ROCM EP since it is not supported yet.
2. For AMD GPU, the EP name is ROCMExecutionProvider
This commit is contained in:
Weixing Zhang 2021-03-10 10:00:35 -08:00 committed by GitHub
parent 8468099f93
commit 1e13e2666e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
1 changed files with 7 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

9
orttraining/orttraining/python/training/ortmodule.py
View file

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