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

[Fix] Fix compile error (#4357)

Browse Source
This commit is contained in:
Mashiro 2023-09-01 18:12:58 +08:00 committed by GitHub
parent eb952ea88d
commit cfa607080f
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
op_builder/utils.py
View File

@ -197,11 +197,12 @@ def get_cuda_cc_flag() -> List[str]:
import torch import torch
cc_flag = [] cc_flag = []
max_arch = ''.join(str(i) for i in torch.cuda.get_device_capability())
for arch in torch.cuda.get_arch_list(): for arch in torch.cuda.get_arch_list():
res = re.search(r'sm_(\d+)', arch) res = re.search(r'sm_(\d+)', arch)
if res: if res:
arch_cap = res[1] arch_cap = res[1]
if int(arch_cap) >= 60: if int(arch_cap) >= 60 and int(arch_cap) <= int(max_arch):
cc_flag.extend(['-gencode', f'arch=compute_{arch_cap},code={arch}']) cc_flag.extend(['-gencode', f'arch=compute_{arch_cap},code={arch}'])
return cc_flag return cc_flag