[Core] manage nccl via a pypi package & upgrade to pt 2.2.1 (#3805)

This commit is contained in:
youkaichao
2024-04-04 10:26:19 -07:00
committed by GitHub
parent b7782002e1
commit ca81ff5196
8 changed files with 36 additions and 11 deletions

View File

@ -328,6 +328,16 @@ def get_requirements() -> List[str]:
if _is_cuda():
with open(get_path("requirements.txt")) as f:
requirements = f.read().strip().split("\n")
cuda_major = torch.version.cuda.split(".")[0]
modified_requirements = []
for req in requirements:
if "vllm-nccl-cu12" in req:
modified_requirements.append(
req.replace("vllm-nccl-cu12",
f"vllm-nccl-cu{cuda_major}"))
else:
modified_requirements.append(req)
requirements = modified_requirements
elif _is_hip():
with open(get_path("requirements-rocm.txt")) as f:
requirements = f.read().strip().split("\n")