[WIP][Core] try to manage nccl version

Dockerfile

@@ -15,26 +15,27 @@ RUN ldconfig /usr/local/cuda-12.1/compat/
 
 WORKDIR /workspace
 
-# install build and runtime dependencies
-COPY requirements.txt requirements.txt
+# install build dependencies
+COPY requirements-build.txt requirements-build.txt
 RUN --mount=type=cache,target=/root/.cache/pip \
-    pip install -r requirements.txt
+    pip install -r requirements-build.txt
 
 # install development dependencies
 COPY requirements-dev.txt requirements-dev.txt
 RUN --mount=type=cache,target=/root/.cache/pip \
     pip install -r requirements-dev.txt
+
+# install runtime dependencies
+COPY requirements.txt requirements.txt
+RUN --mount=type=cache,target=/root/.cache/pip \
+    pip install -r requirements.txt
+
 #################### BASE BUILD IMAGE ####################
 
 
 #################### EXTENSION BUILD IMAGE ####################
 FROM dev AS build
 
-# install build dependencies
-COPY requirements-build.txt requirements-build.txt
-RUN --mount=type=cache,target=/root/.cache/pip \
-    pip install -r requirements-build.txt
-
 # install compiler cache to speed up compilation leveraging local or remote caching
 RUN apt-get update -y && apt-get install -y ccache
 

pyproject.toml

@@ -7,6 +7,7 @@ requires = [
     "setuptools >= 49.4.0",
     "torch == 2.1.2",
     "wheel",
+    "requests", # Required for downloading nccl library
 ]
 build-backend = "setuptools.build_meta"
 

requirements-build.txt

@@ -5,3 +5,4 @@ packaging
 setuptools>=49.4.0
 torch==2.1.2
 wheel
+requests # Required for downloading nccl library

requirements.txt

@@ -17,3 +17,4 @@ pynvml == 11.5.0
 triton >= 2.1.0
 outlines == 0.0.34
 tiktoken == 0.6.0 # Required for DBRX tokenizer
+git+https://github.com/vllm-project/vllm-nccl.git # manage NCCL version to be separate from PyTorch

vllm/model_executor/parallel_utils/pynccl.py

@@ -22,6 +22,8 @@
 import ctypes
 import datetime
 import os
+import site
+import sys
 
 # ===================== import region =====================
 import torch
@@ -40,7 +42,13 @@
         f"Loading nccl from environment variable VLLM_NCCL_SO_PATH={so_file}")
 else:
     if torch.version.cuda is not None:
-        so_file = "libnccl.so.2"
+        # check if we have vllm-managed nccl
+        if os.path.exists(os.path.join(sys.prefix, "vllm_nccl")):
+            so_file = os.path.join(sys.prefix, "vllm_nccl", "libnccl.so.2")
+        elif os.path.exists(os.path.join(site.USER_BASE, "vllm_nccl")):
+            so_file = os.path.join(site.USER_BASE, "vllm_nccl", "libnccl.so.2")
+        else:
+            so_file = "libnccl.so.2"
     elif torch.version.hip is not None:
         so_file = "librccl.so.1"
     else: