[diffusion] hardware: support diffusion models on MTGPU (multi-GPU, 5/N) (#17318)
Signed-off-by: Xiaodong Ye <xiaodong.ye@mthreads.com>
This commit is contained in:
@@ -52,8 +52,8 @@ 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.
|
||||
After importing `torch`, `libnccl.so.2`, `librccl.so.1` or `libmccl.so.2`
|
||||
can be found by `ctypes` automatically.
|
||||
"""
|
||||
so_file = envs.SGLANG_DIFFUSION_NCCL_SO_PATH
|
||||
|
||||
@@ -68,8 +68,10 @@ def find_nccl_library() -> str:
|
||||
so_file = "libnccl.so.2"
|
||||
elif torch.version.hip is not None:
|
||||
so_file = "librccl.so.1"
|
||||
elif hasattr(torch.version, "musa") and torch.version.musa is not None:
|
||||
so_file = "libmccl.so.2"
|
||||
else:
|
||||
raise ValueError("NCCL only supports CUDA and ROCm backends.")
|
||||
raise ValueError("NCCL only supports CUDA, ROCm and MUSA backends.")
|
||||
logger.info("Found nccl from library %s", so_file)
|
||||
return str(so_file)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user