A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 17:15 UTC
[TPU] Lazy Import (#15656)
Signed-off-by: rshaw@neuralmagic.com <robertgshaw2@gmail.com>
Robert Shaw committed
Mar 27, 2025 at 21:57 UTC
bd45912b99e3bad6621fd4d6bc103352ff31bcb7
1 file changed
+2
-1
vllm/distributed/utils.py
+2
-1
@@ -15,7 +15,6 @@ import torch
15
from torch.distributed import ProcessGroup, TCPStore
16
from torch.distributed.distributed_c10d import (Backend, PrefixStore,
17
_get_default_timeout,
18
- _shutdown_backend,
18
_unregister_process_group,
19
is_nccl_available)
20
from torch.distributed.rendezvous import rendezvous
@@ -343,5 +342,7 @@ def stateless_destroy_torch_distributed_process_group(
342
Destroy ProcessGroup returned by
343
stateless_init_torch_distributed_process_group().
344
"""
345
+ # Lazy import for non-CUDA backends.
346
+ from torch.distributed.distributed_c10d import _shutdown_backend
347
_shutdown_backend(pg)
348
_unregister_process_group(pg.group_name)