fix distributed allreduce calling python

norch/distributed/distributed.py

@@ -18,7 +18,7 @@ def broadcast_tensor(tensor):
 
 def allreduce_sum_tensor(tensor):
 
-    Tensor._C.allreduce_mean_tensor.argtypes = [ctypes.POINTER(CTensor)]
-    Tensor._C.allreduce_mean_tensor.restype = None
+    Tensor._C.allreduce_sum_tensor.argtypes = [ctypes.POINTER(CTensor)]
+    Tensor._C.allreduce_sum_tensor.restype = None
     
-    Tensor._C.allreduce_mean_tensor(tensor)
+    Tensor._C.allreduce_sum_tensor(tensor)

norch/distributed/run/run.py

@@ -22,6 +22,7 @@ def main():
     mpiexec_command = [
         'mpiexec',
         '-n', str(world_size),
+        '--allow-run-as-root',
         sys.executable, script_name
     ] + script_args
 

setup.py

@@ -22,7 +22,6 @@ apt_nccl = [
 ]
 
 subprocess.check_call(['sudo', 'apt-get', 'update'])
-subprocess.check_call(['sudo', 'apt-get', 'upgrade', '-y', '--allow-change-held-packages'])
 
 for package in apt_dependencies:
     subprocess.check_call(['sudo', 'apt', 'install', package])
@@ -33,6 +32,8 @@ for package in apt_get_dependencies:
 subprocess.check_call(['wget', 'https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2004/x86_64/cuda-keyring_1.0-1_all.deb'])
 subprocess.check_call(['sudo', 'dpkg', '-i', 'cuda-keyring_1.0-1_all.deb'])
 
+subprocess.check_call(['sudo', 'apt-mark', 'unhold', 'libnccl-dev', 'libnccl2'])
+
 for package in apt_nccl:
     subprocess.check_call(['sudo', 'apt', 'install', package])