github/ColossalAI
1
0
Fork 0
mirror of https://github.com/hpcaitech/ColossalAI.git synced 2025-09-01 01:06:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

[builder] raise Error when CUDA_HOME is not set (#2213)

Browse Source
This commit is contained in:
Jiarui Fang
2022-12-28 16:07:08 +08:00
committed by GitHub
parent 78a89d9b41
commit 7675792100
7 changed files with 19 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
colossalai/kernel/op_builder/builder.py
View File

@@ -30,6 +30,13 @@ class Builder(object):
else:
return os.path.join(Path(__file__).parent.parent.absolute(), code_path)
def get_cuda_include(self):
from torch.utils.cpp_extension import CUDA_HOME
if CUDA_HOME is None:
raise RuntimeError("CUDA_HOME is None, please set CUDA_HOME to compile C++/CUDA kernels in ColossalAI.")
cuda_include = os.path.join(CUDA_HOME, "include")
return cuda_include
def strip_empty_entries(self, args):
'''
Drop any empty strings from the list of compile and link flags

4
colossalai/kernel/op_builder/cpu_adam.py
View File

@@ -27,9 +27,7 @@ class CPUAdamBuilder(Builder):
]
def include_paths(self):
from torch.utils.cpp_extension import CUDA_HOME
cuda_include = os.path.join(CUDA_HOME, "include")
return [os.path.join(CPUAdamBuilder.BASE_DIR, "includes"), cuda_include]
return [os.path.join(CPUAdamBuilder.BASE_DIR, "includes"), self.get_cuda_include()]
def strip_empty_entries(self, args):
'''

5
colossalai/kernel/op_builder/fused_optim.py
View File

@@ -31,10 +31,7 @@ class FusedOptimBuilder(Builder):
]
def include_paths(self):
import torch
from torch.utils.cpp_extension import CUDA_HOME
cuda_include = os.path.join(CUDA_HOME, "include")
return [os.path.join(FusedOptimBuilder.BASE_DIR, "includes"), cuda_include]
return [os.path.join(FusedOptimBuilder.BASE_DIR, "includes"), self.get_cuda_include()]
def builder(self, name):
from torch.utils.cpp_extension import CUDAExtension

4
colossalai/kernel/op_builder/multi_head_attn.py
View File

@@ -31,10 +31,8 @@ class MultiHeadAttnBuilder(Builder):
]
def include_paths(self):
from torch.utils.cpp_extension import CUDA_HOME
ret = []
cuda_include = os.path.join(CUDA_HOME, "include")
ret = [os.path.join(self.base_dir, "includes"), cuda_include]
ret = [os.path.join(self.base_dir, "includes"), self.get_cuda_include()]
ret.append(os.path.join(self.base_dir, "kernels", "include"))
print("include_paths", ret)
return ret