frozenleaves/DeepSpeed
1
0
Fork 0
mirror of https://github.com/deepspeedai/DeepSpeed.git synced 2025-10-20 23:53:48 +08:00
Code Issues Packages Projects Releases Wiki Activity

AIO on ROCM (#7023)

Browse Source 
Adding compile support for AIO library on AMD GPUs.

---------

Co-authored-by: Olatunji Ruwase <olruwase@microsoft.com>
Co-authored-by: Logan Adams <114770087+loadams@users.noreply.github.com>
This commit is contained in:
Joe Mayer
2025-02-14 09:17:08 -08:00
committed by GitHub
parent e637677766
commit 14b3cce4aa
1 changed files with 5 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
op_builder/async_io.py
View File

@ -65,8 +65,11 @@ class AsyncIOBuilder(TorchCPUOpBuilder):
import torch.utils.cpp_extension
CUDA_HOME = torch.utils.cpp_extension.CUDA_HOME
CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64")
ldflags = [f'-L{CUDA_HOME}', f'-L{CUDA_LIB64}', '-laio', '-lcuda', '-lcudart']
if CUDA_HOME is None:
ldflags = ['-laio'] # the ROCM case
else:
CUDA_LIB64 = os.path.join(CUDA_HOME, "lib64")
ldflags = [f'-L{CUDA_HOME}', f'-L{CUDA_LIB64}', '-laio', '-lcuda', '-lcudart']
return ldflags
def check_for_libaio_pkg(self):