Install libnccl-dev in base image since JAX no longer vendors NCCL (#289

)
NVIDIA · Oct 7, 2023 · 003240c · 003240c
1 parent ed713a9
commit 003240c
Showing 1 changed file with 9 additions and 0 deletions.
diff --git a/.github/container/Dockerfile.base b/.github/container/Dockerfile.base
@@ -30,6 +30,15 @@ RUN pip install --upgrade --no-cache-dir pip
 ADD install-cudnn.sh /usr/local/bin
 RUN install-cudnn.sh
 
+###############################################################################
+## Install NCCL
+###############################################################################
+
+RUN apt-get update && \
+    apt-get install -y libnccl-dev && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
 ###############################################################################
 ## RoCE and InfiniteBand support
 ###############################################################################