From 6d8ff7f9af4d0bd00fa46c42c910988eda1c6fc4 Mon Sep 17 00:00:00 2001 From: lucasdelimanogueira Date: Fri, 24 May 2024 20:06:34 -0300 Subject: [PATCH] fix distributed allreduce calling python --- norch/distributed/distributed.py | 6 +++--- norch/distributed/run/run.py | 1 + setup.py | 3 ++- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/norch/distributed/distributed.py b/norch/distributed/distributed.py index d033ede..b078664 100644 --- a/norch/distributed/distributed.py +++ b/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) diff --git a/norch/distributed/run/run.py b/norch/distributed/run/run.py index f40cb63..2e18211 100644 --- a/norch/distributed/run/run.py +++ b/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 diff --git a/setup.py b/setup.py index 0d146ee..e8d6c28 100644 --- a/setup.py +++ b/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])