Enable symm mem in DP setup

Signed-off-by: ilmarkov <markovilya197@gmail.com>

Author: ilmarkov

vllm/distributed/device_communicators/cuda_communicator.py

@@ -24,19 +24,18 @@ def __init__(self,
                  unique_name: str = ""):
         super().__init__(cpu_group, device, device_group, unique_name)
         if "tp" not in unique_name:
-            # only tp uses custom allreduce
+            # custom allreduce or torch symm mem can be used only by tp
             use_custom_allreduce = False
+            use_torch_symm_mem = False
         else:
             from vllm.distributed.parallel_state import (
                 _ENABLE_CUSTOM_ALL_REDUCE)
             use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
+            use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM
 
         # ep does not use pynccl
         use_pynccl = "ep" not in unique_name
 
-        use_torch_symm_mem = ("ep" not in unique_name
-                              and envs.VLLM_ALLREDUCE_USE_SYMM_MEM)
-
         self.use_pynccl = use_pynccl
         self.use_custom_allreduce = use_custom_allreduce
         self.use_torch_symm_mem = use_torch_symm_mem

vllm/distributed/device_communicators/symm_mem.py

@@ -1,6 +1,5 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-import os
 from typing import Optional, Union
 
 import torch
@@ -81,11 +80,6 @@ def __init__(
             self.max_size = SYMM_MEM_ALL_REDUCE_MAX_SIZES[
                 self.device_capability][self.world_size]
 
-        # allow overlapping devices in case of data parallel
-        # if torch symm mem can not initialize multicast ptr will be 0
-        if self.world_size < dist.get_world_size():
-            os.environ["TORCH_SYMM_MEM_ALLOW_OVERLAPPING_DEVICES"] = "1"
-
         self.buffer = torch_symm_mem.empty(
             self.max_size // self.dtype.itemsize,
             device=self.device,