github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

Allow building cuda extension without a device. (#5535) 

Added FORCE_CUDA environment variable support, to enable building extensions where a GPU device is not present but cuda libraries are.
pull/5787/head
Charles Coulombe 2024-06-05 02:26:30 -04:00 committed by GitHub
parent 3f7e3131d9
commit c46e09715c
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
extensions/cuda_extension.py
View File

@ -28,7 +28,9 @@ class _CudaExtension(_CppExtension):
try: try:
import torch import torch
cuda_available = torch.cuda.is_available() # torch.cuda.is_available requires a device to exist, allow building with cuda extension on build nodes without a device
# but where cuda is actually available.
cuda_available = torch.cuda.is_available() or bool(os.environ.get('FORCE_CUDA', 0))
except: except:
cuda_available = False cuda_available = False
return cuda_available return cuda_available