Skip to content

GitLab

Unverified Commit 68bb122e authored by Huazhong Ji's avatar Huazhong Ji Committed by GitHub
Browse files

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


Signed-off-by: default avatarhzji210@gmail.com <hzji210@gmail.com>
parent d9fc8cd9
Show whitespace changes
Inline Side-by-side
Showing with 3 additions and 1 deletion
vllm/distributed/parallel_state.py
View file @ 68bb122e
...@@ -194,9 +194,11 @@ class GroupCoordinator: ...@@ -194,9 +194,11 @@ class GroupCoordinator:
from vllm.platforms import current_platform from vllm.platforms import current_platform
# TODO: fix it for other platforms
if current_platform.is_cuda_alike(): if current_platform.is_cuda_alike():
self.device = torch.device(f"cuda:{local_rank}") 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: else:
self.device = torch.device("cpu") self.device = torch.device("cpu")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment