diff --git a/alpa/device_mesh.py b/alpa/device_mesh.py index e74ab418c..4991f81dc 100644 --- a/alpa/device_mesh.py +++ b/alpa/device_mesh.py @@ -1092,6 +1092,7 @@ def launch_xla_servers(self): "FI_EFA_USE_DEVICE_RDMA": "1", "LD_LIBRARY_PATH": os.environ.get("LD_LIBRARY_PATH", ""), # For libnccl-net.so + "NCCL_PROTO": "simple", }) bundle_index = device_bundle_idx_list[i] diff --git a/benchmark/cupy/profile_communication.py b/benchmark/cupy/profile_communication.py index ecc6bde0c..63197e5a4 100644 --- a/benchmark/cupy/profile_communication.py +++ b/benchmark/cupy/profile_communication.py @@ -246,6 +246,7 @@ def sync(self): "FI_PROVIDER": "efa", "FI_EFA_USE_DEVICE_RDMA": "1", "LD_LIBRARY_PATH": os.environ.get("LD_LIBRARY_PATH", ""), # For libnccl-net.so + "NCCL_PROTO": "simple", } elif args.ib: env_vars = { diff --git a/docs/install.rst b/docs/install.rst index 8cb52fe91..41fdebd57 100644 --- a/docs/install.rst +++ b/docs/install.rst @@ -20,8 +20,8 @@ Regardless of installing from wheels or from source, there are a few prerequisit # Update pip pip3 install --upgrade pip - # Use your own CUDA version. Here cuda-cuda114 means cuda 11.4 - pip3 install cupy-cuda114 + # Install cupy + pip3 install cupy-cuda11x Then, check whether your system already has NCCL installed. @@ -32,6 +32,9 @@ Regardless of installing from wheels or from source, there are a few prerequisit If it prints nothing, then NCCL has already been installed. Otherwise, follow the printed instructions to install NCCL. + .. code:: bash + + python3 -m cupyx.tools.install_library --cuda 11.x --library nccl Methods -------