Xuan-1998andVladimir Aerov
026d5af7f4
Signed-off-by: Xuan Jiang <xuanj@amazon.com> Co-authored-by: Vladimir Aerov <vaerov@amazon.com>
130 lines
4.1 KiB
Python
130 lines
4.1 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from __future__ import annotations
|
|
|
|
import ctypes
|
|
import importlib.util
|
|
import os
|
|
|
|
import torch
|
|
|
|
import vllm.envs as envs
|
|
from vllm.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
def find_nccl_library() -> str:
|
|
"""Return NCCL/RCCL shared library name to load.
|
|
|
|
Uses `VLLM_NCCL_SO_PATH` if set; otherwise chooses by torch backend.
|
|
"""
|
|
so_file = envs.VLLM_NCCL_SO_PATH
|
|
if so_file:
|
|
logger.info(
|
|
"Found nccl from environment variable VLLM_NCCL_SO_PATH=%s", so_file
|
|
)
|
|
else:
|
|
if torch.version.cuda is not None:
|
|
so_file = "libnccl.so.2"
|
|
elif torch.version.hip is not None:
|
|
so_file = "librccl.so.1"
|
|
else:
|
|
raise ValueError("NCCL only supports CUDA and ROCm backends.")
|
|
logger.debug_once("Found nccl from library %s", so_file)
|
|
return so_file
|
|
|
|
|
|
def find_nccl_include_paths() -> list[str] | None:
|
|
"""Return possible include paths containing `nccl.h`.
|
|
|
|
Considers `VLLM_NCCL_INCLUDE_PATH` and the `nvidia-nccl-cuXX` package.
|
|
"""
|
|
paths: list[str] = []
|
|
inc = envs.VLLM_NCCL_INCLUDE_PATH
|
|
if inc and os.path.isdir(inc):
|
|
paths.append(inc)
|
|
|
|
try:
|
|
spec = importlib.util.find_spec("nvidia.nccl")
|
|
if spec and (locs := getattr(spec, "submodule_search_locations", None)):
|
|
for loc in locs:
|
|
inc_dir = os.path.join(loc, "include")
|
|
if os.path.exists(os.path.join(inc_dir, "nccl.h")):
|
|
paths.append(inc_dir)
|
|
except Exception as e:
|
|
logger.debug("Failed to find nccl include path from nvidia.nccl package: %s", e)
|
|
|
|
seen: set[str] = set()
|
|
out: list[str] = []
|
|
for p in paths:
|
|
if p and p not in seen:
|
|
out.append(p)
|
|
seen.add(p)
|
|
return out or None
|
|
|
|
|
|
def find_nccl_library_paths() -> list[str] | None:
|
|
"""Return possible library paths containing `libnccl.so`.
|
|
|
|
Looks inside the `nvidia-nccl-cuXX` pip package.
|
|
"""
|
|
paths: list[str] = []
|
|
try:
|
|
spec = importlib.util.find_spec("nvidia.nccl")
|
|
if spec and (locs := getattr(spec, "submodule_search_locations", None)):
|
|
for loc in locs:
|
|
lib_dir = os.path.join(loc, "lib")
|
|
if os.path.isdir(lib_dir):
|
|
paths.append(lib_dir)
|
|
except Exception as e:
|
|
logger.debug("Failed to find nccl library path from nvidia.nccl package: %s", e)
|
|
return paths or None
|
|
|
|
|
|
def query_nccl_gin_type(group: torch.distributed.ProcessGroup) -> int | None:
|
|
"""Return the GIN type for an initialized group, or ``None`` on failure."""
|
|
from vllm.distributed.device_communicators.pynccl_wrapper import (
|
|
NCCL_COMM_PROPERTIES_LAYOUT_VERSION,
|
|
NCCLLibrary,
|
|
ncclCommProperties,
|
|
)
|
|
|
|
try:
|
|
backend = group._get_backend(torch.device("cuda"))
|
|
# GIN is a property of this initialized communicator, not just the
|
|
# NCCL version. ncclCommQueryProperties requires its ncclComm_t.
|
|
comm_ptr = backend._comm_ptr()
|
|
if comm_ptr == 0:
|
|
return None
|
|
except Exception:
|
|
logger.warning(
|
|
"Failed to extract NCCL comm pointer from process group",
|
|
exc_info=True,
|
|
)
|
|
return None
|
|
|
|
try:
|
|
nccl = NCCLLibrary()
|
|
query_fn = nccl._funcs.get("ncclCommQueryProperties")
|
|
if query_fn is None:
|
|
return None
|
|
|
|
props = ncclCommProperties()
|
|
ctypes.memset(ctypes.addressof(props), 0, ctypes.sizeof(props))
|
|
props.size = ctypes.sizeof(props)
|
|
props.magic = 0xCAFEBEEF
|
|
props.version = min(
|
|
nccl.ncclGetRawVersion(), NCCL_COMM_PROPERTIES_LAYOUT_VERSION
|
|
)
|
|
result = query_fn(ctypes.c_void_p(comm_ptr), ctypes.byref(props))
|
|
except Exception:
|
|
logger.warning("Failed to query NCCL communicator properties", exc_info=True)
|
|
return None
|
|
|
|
if result != 0:
|
|
logger.warning("ncclCommQueryProperties returned error %d", result)
|
|
return None
|
|
return props.ginType
|