setup.py 1.36 KB
Newer Older
1
import setuptools
Rick Ho's avatar
Rick Ho committed
2
3
4
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import os

Rick Ho's avatar
Rick Ho committed
5

Rick Ho's avatar
Rick Ho committed
6
7
8
9
CUDA_HELPER = os.environ.get('CUDA_HELPER', '/usr/local/cuda/samples/common/inc')
cxx_flags = [
        '-I{}'.format(CUDA_HELPER)
        ]
Rick Ho's avatar
Rick Ho committed
10
ext_libs = []
Rick Ho's avatar
Rick Ho committed
11
12
if os.environ.get('USE_NCCL', '0') == '1':
    cxx_flags.append('-DMOE_USE_NCCL')
13
    ext_libs.append('nccl')
Rick Ho's avatar
Rick Ho committed
14

Rick Ho's avatar
Rick Ho committed
15

16
17
if __name__ == '__main__':
    setuptools.setup(
Rick Ho's avatar
Rick Ho committed
18
19
20
        name='fastmoe',
        version='0.1.1',
        description='An efficient Mixture-of-Experts system for PyTorch',
Rick Ho's avatar
Rick Ho committed
21
22
23
24
        author='Jiaao He, Jiezhong Qiu and Aohan Zeng',
        author_email='hja20@mails.tsinghua.edu.cn',
        license='Apache-2',
        url='https://github.com/laekov/fastmoe',
Rick Ho's avatar
Rick Ho committed
25
        packages=['fmoe'],
26
27
28
29
30
31
32
33
34
35
36
37
38
        ext_modules=[
            CUDAExtension(
                name='fmoe_cuda', 
                sources=[
                    'cuda/moe.cpp',
                    'cuda/cuda_stream_manager.cpp',
                    'cuda/moe_compute_kernel.cu',
                    'cuda/moe_comm_kernel.cu',
                    'cuda/moe_fused_kernel.cu',
                    ],
                extra_compile_args={
                    'cxx': cxx_flags,
                    'nvcc': cxx_flags
Rick Ho's avatar
Rick Ho committed
39
40
                    },
                libraries=ext_libs
41
42
43
44
45
                )
            ],
        cmdclass={
            'build_ext': BuildExtension
        })