20231088/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[MISC] Make GroupCoordinator compatible with out-of-tree devices (#16464)

Browse Source 
Signed-off-by: hzji210@gmail.com <hzji210@gmail.com>
This commit is contained in:
Huazhong Ji 2025-04-12 17:20:25 +08:00 committed by GitHub
parent d9fc8cd9da
commit 68bb122eb4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
vllm/distributed/parallel_state.py
View File

@ -194,9 +194,11 @@ class GroupCoordinator:
from vllm.platforms import current_platform
# TODO: fix it for other platforms
if current_platform.is_cuda_alike():
self.device = torch.device(f"cuda:{local_rank}")
elif current_platform.is_out_of_tree():
self.device = torch.device(
f"{current_platform.device_name}:{local_rank}")
else:
self.device = torch.device("cpu")