use custom_all_reduce when bs*hidden_size<=256*4096 (bw)

vllm/distributed/device_communicators/custom_all_reduce.py

@@ -54,7 +54,7 @@ class CustomAllreduce:
     def __init__(self,
                  group: ProcessGroup,
                  device: Union[int, str, torch.device],
-                 max_size=8192 * 1024) -> None:
+                 max_size=8192 * 512) -> None:
         """
         Args:
             group: the process group to work on. If None, it will use the
@@ -230,7 +230,7 @@ class CustomAllreduce:
             return False
         # for 4 or more non NVLink-capable GPUs, custom allreduce provides
         # little performance improvement over NCCL.
-        return inp_size < self.max_size
+        return inp_size <= self.max_size
 
     def all_reduce(self,
                    inp: torch.Tensor,