fix

Signed-off-by: Yongji Wu <[email protected]>

Author: libertyeagle

vllm/distributed/stateless_coordinator.py

@@ -211,9 +211,9 @@ def broadcast_tensor_dict(
             if tensor.numel() == 0:
                 continue
             if self.device_communicator and tensor.is_cuda:
-                self.device_communicator.broadcast(tensor, src)
+                tensor.copy_(self.device_communicator.broadcast(tensor, src))
             else:
-                self.tcp_store_group.broadcast(tensor, src)
+                tensor.copy_(self.tcp_store_group.broadcast(tensor, src))
 
         return tensor_dict
 

vllm/distributed/utils.py

@@ -15,7 +15,7 @@
 from collections import deque
 from collections.abc import Sequence
 from datetime import timedelta
-from typing import Any, Optional
+from typing import Any, Optional, Union
 
 import torch
 from torch.distributed import ProcessGroup, Store, TCPStore
@@ -536,7 +536,7 @@ def stateless_init_torch_distributed_process_group(
     backend: str,
     group_name: Optional[str] = None,
     return_store: bool = False,
-) -> ProcessGroup | tuple[ProcessGroup, Store]:
+) -> Union[ProcessGroup, tuple[ProcessGroup, Store]]:
     """
     A replacement for `torch.distributed.init_process_group` that does not
     pollute the global state. The created ProcessGroup object can be used for