From bd45912b99e3bad6621fd4d6bc103352ff31bcb7 Mon Sep 17 00:00:00 2001 From: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> Date: Thu, 27 Mar 2025 21:57:01 -0400 Subject: [PATCH] [TPU] Lazy Import (#15656) Signed-off-by: rshaw@neuralmagic.com --- vllm/distributed/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/distributed/utils.py b/vllm/distributed/utils.py index b8178af5a2daa..4206a24465e28 100644 --- a/vllm/distributed/utils.py +++ b/vllm/distributed/utils.py @@ -15,7 +15,6 @@ import torch from torch.distributed import ProcessGroup, TCPStore from torch.distributed.distributed_c10d import (Backend, PrefixStore, _get_default_timeout, - _shutdown_backend, _unregister_process_group, is_nccl_available) from torch.distributed.rendezvous import rendezvous @@ -343,5 +342,7 @@ def stateless_destroy_torch_distributed_process_group( Destroy ProcessGroup returned by stateless_init_torch_distributed_process_group(). """ + # Lazy import for non-CUDA backends. + from torch.distributed.distributed_c10d import _shutdown_backend _shutdown_backend(pg) _unregister_process_group(pg.group_name)