forked from Karylab-cklius/vllm
Address review: add push_ar to benchmark and backend logging
- Add PushAllReduce to benchmark_device_communicators.py for comparing against other allreduce implementations - Add PUSH_AR to _log_all_reduce_backend_selection in cuda_communicator.py for visibility in dispatch logging Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: Alexander Matveev <alexm-redhat@dgx-b200-02.mgmt.accl-001.lab.rdu2.dc.redhat.com>
This commit is contained in:
co-authored by
Claude Opus 4.6
parent
bc50e5fc2e
commit
fd44100bb0
@@ -33,6 +33,7 @@ from vllm.distributed.device_communicators.custom_all_reduce import CustomAllred
|
||||
from vllm.distributed.device_communicators.flashinfer_all_reduce import (
|
||||
FlashInferAllReduce,
|
||||
)
|
||||
from vllm.distributed.device_communicators.push_all_reduce import PushAllReduce
|
||||
from vllm.distributed.device_communicators.pynccl import (
|
||||
PyNcclCommunicator,
|
||||
register_nccl_symmetric_ops,
|
||||
@@ -80,6 +81,7 @@ class CommunicatorBenchmark:
|
||||
|
||||
# Initialize communicators
|
||||
self.custom_allreduce = None
|
||||
self.push_ar_comm = None
|
||||
self.pynccl_comm = None
|
||||
self.symm_mem_comm = None
|
||||
self.symm_mem_comm_multimem = None
|
||||
@@ -106,6 +108,23 @@ class CommunicatorBenchmark:
|
||||
)
|
||||
self.custom_allreduce = None
|
||||
|
||||
try:
|
||||
self.push_ar_comm = PushAllReduce(
|
||||
group=self.cpu_group,
|
||||
device=self.device,
|
||||
max_size=self.max_size_override,
|
||||
)
|
||||
if not self.push_ar_comm.disabled:
|
||||
logger.info("Rank %s: PushAllReduce initialized", self.rank)
|
||||
else:
|
||||
logger.info("Rank %s: PushAllReduce disabled", self.rank)
|
||||
self.push_ar_comm = None
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"Rank %s: Failed to initialize PushAllReduce: %s", self.rank, e
|
||||
)
|
||||
self.push_ar_comm = None
|
||||
|
||||
try:
|
||||
self.pynccl_comm = PyNcclCommunicator(
|
||||
group=self.cpu_group, device=self.device
|
||||
@@ -216,6 +235,19 @@ class CommunicatorBenchmark:
|
||||
)
|
||||
)
|
||||
|
||||
if self.push_ar_comm is not None:
|
||||
comm = self.push_ar_comm
|
||||
communicators.append(
|
||||
(
|
||||
"push_ar",
|
||||
lambda t, c=comm: c.all_reduce(t),
|
||||
lambda t, c=comm: c.should_use(t),
|
||||
comm.capture(),
|
||||
{},
|
||||
None,
|
||||
)
|
||||
)
|
||||
|
||||
if self.pynccl_comm is not None:
|
||||
comm = self.pynccl_comm
|
||||
communicators.append(
|
||||
|
||||
@@ -224,6 +224,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
|
||||
"NCCL_SYMM_MEM",
|
||||
"QUICK_REDUCE",
|
||||
"FLASHINFER",
|
||||
"PUSH_AR",
|
||||
"CUSTOM",
|
||||
"SYMM_MEM",
|
||||
"PYNCCL",
|
||||
@@ -257,6 +258,8 @@ class CudaCommunicator(DeviceCommunicatorBase):
|
||||
enabled_ar_backends.append("QUICK_REDUCE")
|
||||
if self.fi_ar_comm is not None and not self.fi_ar_comm.disabled:
|
||||
enabled_ar_backends.append("FLASHINFER")
|
||||
if self.push_ar_comm is not None:
|
||||
enabled_ar_backends.append("PUSH_AR")
|
||||
if self.ca_comm is not None and not self.ca_comm.disabled:
|
||||
enabled_ar_backends.append("CUSTOM")
|
||||
if self.symm_mem_comm is not None and not self.symm_mem_comm.disabled:
|
||||
|
||||
Reference in New Issue
Block a user