fix nccl for docker cu11 (#3896)

Author: RunningLeon

docker/Dockerfile

@@ -14,6 +14,7 @@ ENV CUDA_VERSION_SHORT=cu118
 # Builder image
 FROM ${CUDA_VERSION} AS dev
 ARG PYTHON_VERSION=3.10
+ARG NCCL_BRANCH=v2.26.6-1
 
 ENV PATH=/opt/py3/bin:/root/.local/bin:${PATH}
 ENV DEBIAN_FRONTEND=noninteractive
@@ -58,11 +59,12 @@ COPY docker/install.sh /tmp/install.sh
 
 RUN --mount=type=cache,target=/root/.cache \
     --mount=type=cache,target=/wheels,from=builder,source=/wheels \
+    --mount=type=cache,target=/nccl,from=builder,source=/nccl \
     /tmp/install.sh
 
 # explicitly set ptxas path for triton
 ENV PATH=/opt/py3/bin:$PATH
 ENV TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas
 ENV NCCL_LAUNCH_MODE=GROUP
-
+ENV LD_LIBRARY_PATH=/nccl/lib:$LD_LIBRARY_PATH
 FROM ${IMAGE_TYPE}

docker/build.sh

@@ -1,11 +1,23 @@
 #!/bin/bash -ex
 
-mkdir -p /wheels
+mkdir -p /wheels /nccl
 
 if [[ "${CUDA_VERSION_SHORT}" != "cu118" ]]; then
     pip install nvidia-nccl-cu12
 else
-    pip install nvidia-nccl-cu11
+    NVCC_GENCODE="-gencode=arch=compute_70,code=sm_70 -gencode=arch=compute_75,code=sm_75 -gencode=arch=compute_80,code=sm_80 -gencode=arch=compute_90,code=sm_90 -gencode=arch=compute_90,code=compute_90"
+    pushd /tmp >/dev/null
+        git clone --depth=1 --branch ${NCCL_BRANCH} https://github.com/NVIDIA/nccl.git
+        pushd nccl >/dev/null
+            make NVCC_GENCODE="$NVCC_GENCODE" -j$(nproc) src.build
+            mkdir -p /nccl/include /nccl/lib
+            mv build/include/* /nccl/include/
+            mv build/lib/lib* /nccl/lib/
+        popd >/dev/null
+    popd >/dev/null
+    rm -rf /tmp/nccl
+    export LD_LIBRARY_PATH=/nccl/lib:$LD_LIBRARY_PATH
 fi
 
+pip install --upgrade pip build
 python3 -m build -w -o /wheels -v .

docker/install.sh

@@ -70,3 +70,9 @@ BASE_URL="https://github.com/Dao-AILab/flash-attention/releases/download/v${FA_V
 FULL_URL="${BASE_URL}/${WHEEL}"
 
 pip install "$FULL_URL"
+
+# copy nccl
+if [[ "${CUDA_VERSION_SHORT}" = "cu118" ]]; then
+    rm -rf /opt/py3/lib/python${PYTHON_VERSION}/site-packages/nvidia/nccl
+    cp -R /nccl /opt/py3/lib/python${PYTHON_VERSION}/site-packages/nvidia/
+fi