refactor reduce

Author: shihaobai

lightllm/distributed/communication_op.py

@@ -24,6 +24,7 @@
 import torch.distributed as dist
 from torch.distributed import ReduceOp
 from lightllm.utils.log_utils import init_logger
+from functools import partial
 
 original_all_reduce = torch.distributed.all_reduce
 from contextlib import nullcontext, contextmanager
@@ -47,31 +48,30 @@ def lightllm_capture_graph():
     pass
 
 def _all_reduce(input_, op=ReduceOp.SUM, group=None, async_op=False):
-    if op != ReduceOp.SUM or group is not None or async_op:
+    if op != ReduceOp.SUM or async_op:
         original_all_reduce(input_, op, group, async_op)
     else:
         if vllm_reduce is not None:
             can_use = vllm_reduce.should_custom_ar(input_)
             if can_use:
                 input_.data = vllm_reduce.custom_all_reduce(input_)
                 return
-            original_all_reduce(input_, op, vllm_reduce.device_group, async_op)
-        else:
-            original_all_reduce(input_, op, group, async_op)
-
+        original_all_reduce(input_, op, group, async_op)
 
 def set_custom_reduce():
     global vllm_reduce
+    global device_group
     ENABLE_VLLM_REDUCE = os.getenv("ENABLE_VLLM_REDUCE", "False").upper() in [
         "ON",
         "TRUE",
         "1",
     ]
+    world_size = dist.get_world_size()
+    ranks = list(range(world_size))
+    # new_group prevent stuck of torch origin all_reduce with cudagraph
+    device_group = torch.distributed.new_group(ranks, backend="nccl")
     if ENABLE_VLLM_REDUCE and HAS_VLLM:
-        world_size = dist.get_world_size()
-        ranks = list(range(world_size))
-        device_group = torch.distributed.new_group(ranks, backend="nccl")
         cpu_group = torch.distributed.new_group(ranks, backend="gloo")
-        vllm_reduce = CustomAllreduce(cpu_group, device_group, torch.cuda.current_device())
+        vllm_reduce = CustomAllreduce(cpu_group, torch.cuda.current_device())
         logger.info("Enable VLLM ALLReduce.")
-        dist.all_reduce = _all_reduce
+    dist.all_reduce = partial(_all_reduce, group=device_group)

lightllm/distributed/custom_all_reduce.py

@@ -49,7 +49,7 @@ class CustomAllreduce:
     _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8]
 
     # max_size: max supported allreduce size
-    def __init__(self, group: ProcessGroup, device_group: ProcessGroup, device: Union[int, str, torch.device], max_size=8192 * 1024) -> None:
+    def __init__(self, group: ProcessGroup, device: Union[int, str, torch.device], max_size=8192 * 1024) -> None:
         """
         Args:
             group: the process group to work on. If None, it will use the
@@ -69,7 +69,6 @@ def __init__(self, group: ProcessGroup, device_group: ProcessGroup, device: Unio
             return
 
         self.group = group
-        self.device_group = device_group
         assert dist.get_backend(group) != dist.Backend.NCCL, "CustomAllreduce should be attached to a non-NCCL group."
 
         rank = dist.get_rank(group=self.group)