66 lines
2.2 KiB
Python
66 lines
2.2 KiB
Python
import torch
|
|
import fastvideo.envs as envs
|
|
from fastvideo.logger import init_logger
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
def find_nccl_library() -> str:
|
|
"""
|
|
We either use the library file specified by the `VLLM_NCCL_SO_PATH`
|
|
environment variable, or we find the library file brought by PyTorch.
|
|
After importing `torch`, `libnccl.so.2` or `librccl.so.1` can be
|
|
found by `ctypes` automatically.
|
|
"""
|
|
so_file = envs.FASTVIDEO_NCCL_SO_PATH
|
|
|
|
# manually load the nccl library
|
|
if so_file:
|
|
logger.info(
|
|
"Found nccl from environment variable FASTVIDEO_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.info("Found nccl from library %s", so_file)
|
|
return so_file
|
|
|
|
prev_set_stream = torch.cuda.set_stream
|
|
|
|
_current_stream = None
|
|
|
|
|
|
def _patched_set_stream(stream: torch.cuda.Stream) -> None:
|
|
global _current_stream
|
|
_current_stream = stream
|
|
prev_set_stream(stream)
|
|
|
|
|
|
torch.cuda.set_stream = _patched_set_stream
|
|
|
|
|
|
def current_stream() -> torch.cuda.Stream:
|
|
"""
|
|
replace `torch.cuda.current_stream()` with `vllm.utils.current_stream()`.
|
|
it turns out that `torch.cuda.current_stream()` is quite expensive,
|
|
as it will construct a new stream object at each call.
|
|
here we patch `torch.cuda.set_stream` to keep track of the current stream
|
|
directly, so that we can avoid calling `torch.cuda.current_stream()`.
|
|
|
|
the underlying hypothesis is that we do not call `torch._C._cuda_setStream`
|
|
from C/C++ code.
|
|
"""
|
|
from fastvideo.platforms import current_platform
|
|
global _current_stream
|
|
if _current_stream is None:
|
|
# when this function is called before any stream is set,
|
|
# we return the default stream.
|
|
# On ROCm using the default 0 stream in combination with RCCL
|
|
# is hurting performance. Therefore creating a dedicated stream
|
|
# per process
|
|
_current_stream = torch.cuda.Stream() if current_platform.is_rocm(
|
|
) else torch.cuda.current_stream()
|
|
return _current_stream |