Skip to content

GitLab

"git@developer.sourcefind.cn:OpenDAS/torchaudio.git" did not exist on "30c7077b5119664a23b0d1b266ff6856a7c3d818"
Commit 3501fa41 authored by Rick Ho's avatar Rick Ho
Browse files

add nccl flag in setup.py

parent dbf55073
Hide whitespace changes
Inline Side-by-side
Showing with 7 additions and 8 deletions
pytorch/cuda/setup.py
View file @ 3501fa41
...@@ -3,6 +3,11 @@ from torch.utils.cpp_extension import BuildExtension, CUDAExtension ...@@ -3,6 +3,11 @@ from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import os import os
CUDA_HELPER = os.environ.get('CUDA_HELPER', '/usr/local/cuda/samples/common/inc') CUDA_HELPER = os.environ.get('CUDA_HELPER', '/usr/local/cuda/samples/common/inc')
cxx_flags = [
'-I{}'.format(CUDA_HELPER)
]
if os.environ.get('USE_NCCL', '0') == '1':
cxx_flags.append('-DMOE_USE_NCCL')
setup( setup(
name='moe_cuda', name='moe_cuda',
...@@ -15,14 +20,8 @@ setup( ...@@ -15,14 +20,8 @@ setup(
'moe_cuda_kernel.cu', 'moe_cuda_kernel.cu',
], ],
extra_compile_args={ extra_compile_args={
'cxx': [ 'cxx': cxx_flags,
'-I{}'.format(CUDA_HELPER), 'nvcc': cxx_flags
'-DMOE_USE_NCCL'
],
'nvcc': [
'-I{}'.format(CUDA_HELPER),
'-DMOE_USE_NCCL'
]
} }
) )
], ],
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment