fix mpi linking error when using nccl

README.md

@@ -8,9 +8,19 @@ model for PyTorch.
 
 ## Installation
 
-PyTorch with CUDA is supported. The repository is currently tested with PyTorch
+### Prerequisites
+
+PyTorch with CUDA is required. The repository is currently tested with PyTorch
 v1.6.0 and CUDA 10, with designed compatibility to other versions.
 
+If distributed version is enabled, NCCL with P2P communication support,
+typically versions >= 2.7.5 is needed. Note that the MPI backend is used as
+there are some necessary messages to be passed by MPI in FMoE, the backend
+should be `mpi`. However, as there are other data to be synchronized by
+`torch.distributed`, cuda-aware mpi is required.
+
+### Installing
+
 Fast MoE contains a set of PyTorch customized opearators, including both C and
 Python components. Use `python setup.py install` to easily install and enjoy
 using Fast MoE for training.

setup.py

-from setuptools import setup
+import setuptools
 from torch.utils.cpp_extension import BuildExtension, CUDAExtension
 import os
 
@@ -8,25 +8,28 @@ cxx_flags = [
         ]
 if os.environ.get('USE_NCCL', '0') == '1':
     cxx_flags.append('-DMOE_USE_NCCL')
+    os.environ['CXX'] = 'mpicxx'
 
-setup(
-    name='moe_cuda',
-    ext_modules=[
-        CUDAExtension(
-            name='moe_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
-                }
-            )
-        ],
-    cmdclass={
-        'build_ext': BuildExtension
-    })
+if __name__ == '__main__':
+    setuptools.setup(
+        name='fmoe_cuda',
+        packages=setuptools.find_packages(),
+        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
+                    }
+                )
+            ],
+        cmdclass={
+            'build_ext': BuildExtension
+        })