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

6
cxx_flags = []
Rick Ho's avatar
Rick Ho committed
7
ext_libs = []
8

Rick Ho's avatar
Rick Ho committed
9
10
11
12
13
14
15
authors = [
        'Jiaao He', 
        'Jiezhong Qiu', 
        'Aohan Zeng', 
        'Tiago Antunes', 
        'Jinjun Peng', 
        'Qin Li',
Rick Ho's avatar
Rick Ho committed
16
        'Mingshu Zhai'
Rick Ho's avatar
Rick Ho committed
17
18
]

hclearner's avatar
hclearner committed
19
20
21
22
23
is_rocm_pytorch = False
if torch.__version__ >= '1.5':
    from torch.utils.cpp_extension import ROCM_HOME
    is_rocm_pytorch = True if ((torch.version.hip is not None) and (ROCM_HOME is not None)) else False

Rick Ho's avatar
Rick Ho committed
24
if os.environ.get('USE_NCCL', '1') == '1':
Rick Ho's avatar
Rick Ho committed
25
    cxx_flags.append('-DFMOE_USE_NCCL')
26
    cxx_flags.append('-DUSE_C10D_NCCL')
hclearner's avatar
hclearner committed
27
    if is_rocm_pytorch:
huchen1's avatar
huchen1 committed
28
29
30
31
        ext_libs.append('rccl')
    else:
        ext_libs.append('nccl')

hclearner's avatar
hclearner committed
32
33
if is_rocm_pytorch:
    define_macros=[('FMOE_USE_HIP', None)]
huchen1's avatar
huchen1 committed
34
35
else:
    define_macros=[]
Rick Ho's avatar
Rick Ho committed
36

Rick Ho's avatar
Rick Ho committed
37

38
39
if __name__ == '__main__':
    setuptools.setup(
Rick Ho's avatar
Rick Ho committed
40
        name='fastmoe',
Rick Ho's avatar
Rick Ho committed
41
        version='1.0.0',
Rick Ho's avatar
Rick Ho committed
42
        description='An efficient Mixture-of-Experts system for PyTorch',
Rick Ho's avatar
Rick Ho committed
43
        author=', '.join(authors),
Rick Ho's avatar
Rick Ho committed
44
45
46
        author_email='hja20@mails.tsinghua.edu.cn',
        license='Apache-2',
        url='https://github.com/laekov/fastmoe',
Rick Ho's avatar
Rick Ho committed
47
        packages=['fmoe', 'fmoe.megatron', 'fmoe.gates', 'fmoe.fastermoe'],
48
49
50
51
        ext_modules=[
            CUDAExtension(
                name='fmoe_cuda', 
                sources=[
Rick Ho's avatar
Rick Ho committed
52
53
                    'cuda/stream_manager.cpp',
                    'cuda/local_exchange.cu',
Rick Ho's avatar
Rick Ho committed
54
                    'cuda/balancing.cu',
Rick Ho's avatar
Rick Ho committed
55
                    'cuda/global_exchange.cpp',
56
                    'cuda/parallel_linear.cu',
Rick Ho's avatar
Rick Ho committed
57
                    'cuda/fmoe_cuda.cpp',
Rick Ho's avatar
Rick Ho committed
58
                    'cuda/fastermoe/smart_schedule.cpp',
59
                    ],
huchen1's avatar
huchen1 committed
60
                define_macros=define_macros,
61
62
63
                extra_compile_args={
                    'cxx': cxx_flags,
                    'nvcc': cxx_flags
Rick Ho's avatar
Rick Ho committed
64
65
                    },
                libraries=ext_libs
66
67
68
69
70
                )
            ],
        cmdclass={
            'build_ext': BuildExtension
        })