[Bugfix] Deterministic MoE combine (reduce_scatterv) under VLLM_BATCH_INVARIANT (#45683)
Signed-off-by: shijuzhao <shijuzhao@tencent.com> Co-authored-by: shijuzhao <shijuzhao@tencent.com>
shijuzhao committed
Aug 21, 2026 at 21:41 UTC
cda3868f5dd83f413be978148ad8bb63c1f57c70
2 files changed
+68
-1
vllm/distributed/device_communicators/cuda_communicator.py
+8
-1
@@ -451,7 +451,14 @@ class CudaCommunicator(DeviceCommunicatorBase):
451
output = torch.empty(
452
output_shape, dtype=input_tensor.dtype, device=input_tensor.device
453
)
454
- if sizes is not None and sizes.count(sizes[0]) != len(sizes):
454
+ use_deterministic_rs = envs.VLLM_BATCH_INVARIANT and world_size > 2
455
+ if use_deterministic_rs:
456
+ # Reduce to a fixed root (0) for determinism
457
+ reduced = torch.empty_like(input_tensor)
458
+ sizes = sizes if sizes else [chunk_size] * world_size
459
+ pynccl_comm.reduce(reduced, input_tensor, root=0)
460
+ pynccl_comm.scatter(output, reduced, sizes, root=0)
461
+ elif sizes is not None and sizes.count(sizes[0]) != len(sizes):
462
pynccl_comm.reduce_scatterv(output, input_tensor, sizes=sizes)
463
else:
464
pynccl_comm.reduce_scatter(output, input_tensor)
vllm/distributed/device_communicators/pynccl.py
+60
@@ -319,6 +319,66 @@ class PyNcclCommunicator:
319
split_offset += split_size
320
self.nccl.ncclGroupEnd()
321
322
+ def reduce(
323
+ self,
324
+ output_tensor: torch.Tensor,
325
+ input_tensor: torch.Tensor,
326
+ root: int,
327
+ op: ReduceOp = ReduceOp.SUM,
328
+ stream=None,
329
+ ):
330
+ if self.disabled:
331
+ return
332
+ assert input_tensor.device == self.device, (
333
+ f"this nccl communicator is created to work on {self.device}, "
334
+ f"but the input tensor is on {input_tensor.device}"
335
+ )
336
+ if stream is None:
337
+ stream = current_stream()
338
+ self.nccl.ncclReduce(
339
+ buffer_type(input_tensor.data_ptr()),
340
+ buffer_type(output_tensor.data_ptr()),
341
+ input_tensor.numel(),
342
+ ncclDataTypeEnum.from_torch(input_tensor.dtype),
343
+ ncclRedOpTypeEnum.from_torch(op),
344
+ root,
345
+ self.comm,
346
+ cudaStream_t(stream.cuda_stream),
347
+ )
348
+
349
+ def scatter(
350
+ self,
351
+ output_tensor: torch.Tensor,
352
+ input_tensor: torch.Tensor,
353
+ sizes: list[int],
354
+ root: int = 0,
355
+ stream=None,
356
+ ):
357
+ if self.disabled:
358
+ return
359
+ assert output_tensor.device == self.device, (
360
+ f"this nccl communicator is created to work on {self.device}, "
361
+ f"but the output tensor is on {output_tensor.device}"
362
+ )
363
+ if stream is None:
364
+ stream = current_stream()
365
+ self.nccl.ncclGroupStart()
366
+ if self.rank == root:
367
+ split_offset = 0
368
+ for dst, split_size in enumerate(sizes):
369
+ if split_size == 0:
370
+ continue
371
+
372
+ chunk = input_tensor[split_offset : split_offset + split_size, ...]
373
+ if dst == root:
374
+ output_tensor.copy_(chunk)
375
+ else:
376
+ self.send(chunk, dst, stream)
377
+ split_offset += split_size
378
+ elif sizes[self.rank] > 0:
379
+ self.recv(output_tensor, root, stream)
380
+ self.nccl.ncclGroupEnd()
381
+
382
def send(self, tensor: torch.Tensor, dst: int, stream=None):
383
if self.disabled:
384
return