fix lazy hang

Signed-off-by: Wuxun Zhang <[email protected]>

Author: wuxun-zhang

vllm_gaudi/distributed/device_communicators/hpu_communicator.py

@@ -1,35 +1,51 @@
 # SPDX-License-Identifier: Apache-2.0
 
+from typing import Optional
 import torch
 import torch.distributed as dist
+from torch.distributed import ProcessGroup
 
 from vllm.distributed.device_communicators.base_device_communicator \
     import DeviceCommunicatorBase
-from vllm.distributed.parallel_state import get_dp_group
 from vllm.forward_context import get_forward_context
+from vllm.distributed.parallel_state import get_dp_group
 
 import habana_frameworks.torch as htorch  # noqa: F401
 
 
-def naive_multicast(x: torch.Tensor,
-                    cu_tokens_across_dp_cpu: torch.Tensor) -> torch.Tensor:
-    assert x.dim() == 2, "Input tensor must be 2D"
-    dp_rank = get_dp_group().rank_in_group
-    dp_world_size = get_dp_group().world_size
-    buffer = torch.empty((cu_tokens_across_dp_cpu[-1], x.size(1)),
-                         device=x.device,
-                         dtype=x.dtype)
-    start = 0 if dp_rank == 0 else cu_tokens_across_dp_cpu[dp_rank - 1]
-    end = cu_tokens_across_dp_cpu[dp_rank]
-    buffer[start:end, :].copy_(x)
-    for idx in range(dp_world_size):
-        start = 0 if idx == 0 else cu_tokens_across_dp_cpu[idx - 1]
-        end = cu_tokens_across_dp_cpu[idx]
-        get_dp_group().broadcast(buffer[start:end, :], idx)
-    return buffer
+class HpuCommunicator(DeviceCommunicatorBase):
 
+    def __init__(self,
+                 cpu_group: ProcessGroup,
+                 device: Optional[torch.device] = None,
+                 device_group: Optional[ProcessGroup] = None,
+                 unique_name: str = ""):
+        super().__init__(cpu_group, device, device_group, unique_name)
 
-class HpuCommunicator(DeviceCommunicatorBase):
+        self.dp_group = None
+        self.dp_rank = 0
+        self.dp_world_size = 1
+        # assume EP is enabled along with DP
+        if "ep" in unique_name:
+            self.dp_group = get_dp_group()
+            self.dp_rank = self.dp_group.rank_in_group
+            self.dp_world_size = self.dp_group.world_size
+
+    def naive_multicast(self, x: torch.Tensor,
+                        cu_tokens_across_dp_cpu: torch.Tensor) -> torch.Tensor:
+        assert x.dim() == 2, "Input tensor must be 2D"
+        buffer = torch.empty((cu_tokens_across_dp_cpu[-1], x.size(1)),
+                             device=x.device,
+                             dtype=x.dtype)
+        start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[
+            self.dp_rank - 1]
+        end = cu_tokens_across_dp_cpu[self.dp_rank]
+        buffer[start:end, :].copy_(x)
+        for idx in range(self.dp_world_size):
+            start = 0 if idx == 0 else cu_tokens_across_dp_cpu[idx - 1]
+            end = cu_tokens_across_dp_cpu[idx]
+            self.dp_group.broadcast(buffer[start:end, :], idx)
+        return buffer
 
     def all_reduce(self, input_: torch.Tensor) -> torch.Tensor:
         # FIXME(kzawora): this is a workaround for a bug in Habana PT bridge
@@ -67,19 +83,57 @@ def dispatch(
             router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
         cu_tokens_across_dp_cpu = get_forward_context(
         ).dp_metadata.cu_tokens_across_dp_cpu
-        hidden_states_across_dp = naive_multicast(hidden_states,
-                                                  cu_tokens_across_dp_cpu)
-        router_logits_across_dp = naive_multicast(router_logits,
-                                                  cu_tokens_across_dp_cpu)
+        hidden_states_across_dp = self.naive_multicast(
+            hidden_states, cu_tokens_across_dp_cpu)
+        router_logits_across_dp = self.naive_multicast(
+            router_logits, cu_tokens_across_dp_cpu)
         return hidden_states_across_dp, router_logits_across_dp
 
+    # def dispatch(
+    #         self, hidden_states: torch.Tensor,
+    #         router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+    #     world_size = self.dp_group.world_size
+    #     input_size = hidden_states.size()
+    #     # Allocate output tensor.
+    #     output_size = list(input_size)
+    #     output_size[0] *= world_size
+    #     hidden_states_across_dp = torch.empty(output_size,
+    #                                 dtype=hidden_states.dtype,
+    #                                 device=hidden_states.device)
+    #     # All-gather.
+    #     torch.distributed.all_gather_into_tensor(
+    #         hidden_states_across_dp, hidden_states, group=self.dp_group.device_group)
+
+    #     router_logits_size = router_logits.size()
+    #     router_logits_output_size = list(router_logits_size)
+    #     router_logits_output_size[0] *= world_size
+    #     router_logits_across_dp = torch.empty(router_logits_output_size,
+    #                                 dtype=router_logits.dtype,
+    #                                 device=router_logits.device)
+    #     # All-gather.
+    #     torch.distributed.all_gather_into_tensor(
+    #         router_logits_across_dp, router_logits, group=self.dp_group.device_group)
+    #     return hidden_states_across_dp, router_logits_across_dp
+
     def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
-        dp_rank = get_dp_group().rank_in_group
+        if htorch.utils.internal.is_lazy():
+            htorch.core.mark_step()
         cu_tokens_across_dp_cpu = get_forward_context(
         ).dp_metadata.cu_tokens_across_dp_cpu
-        start = 0 if dp_rank == 0 else cu_tokens_across_dp_cpu[dp_rank - 1]
-        end = cu_tokens_across_dp_cpu[dp_rank]
 
-        all_hidden_states = get_dp_group().all_reduce(hidden_states)
+        start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[
+            self.dp_rank - 1]
+        end = cu_tokens_across_dp_cpu[self.dp_rank]
+
+        all_hidden_states = self.dp_group.all_reduce(hidden_states)
         hidden_states = all_hidden_states[start:end, :]
         return hidden_states
+
+    # def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
+    #     if htorch.utils.internal.is_lazy():
+    #         htorch.core.mark_step()
+
+    #     all_hidden_states = self.dp_group.all_reduce(hidden_states)
+    #     all_hidden_states = all_hidden_states.view(self.dp_group.world_size, -1, all_hidden_states.size(-1))
+    #     hidden_states = all_hidden_states[self.dp_rank // self.dp_world_size, :, :]
+    #     return hidden_states