add_twomicrobatch

Author: shenhao

lmdeploy/pytorch/backends/cuda/moe.py

@@ -1,6 +1,6 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 
-from typing import List
+from typing import List, Any
 
 import torch
 import torch.distributed as dist
@@ -361,6 +361,15 @@ def __init__(self,
                 hidden_size=hidden_dim,
                 params_dtype=out_dtype,
             )
+        self.token_dispatcher_for2mb = DeepEPDispatcher(
+                group=ep_group,
+                router_topk=self.top_k,
+                permute_fusion=True,
+                num_experts=self.num_experts,
+                num_local_experts=self.num_experts // ep_size,
+                hidden_size=hidden_dim,
+                params_dtype=out_dtype,
+            )
         self.experts = DeepEPMoE(num_experts, ep_size, [block_size,block_size])
 
     def forward(self,
@@ -377,7 +386,7 @@ def forward(self,
         recv_hidden_states, recv_topk_ids, recv_topk_weights, tokens_per_expert = (
             self.token_dispatcher.dispatch(
                 hidden_states,
-                topk_ids.to(torch.int32),
+                topk_ids.to(torch.int64),
                 topk_weights.to(torch.float32),
                 self.num_experts,
             )
@@ -386,6 +395,37 @@ def forward(self,
                                  down_weights, down_scale)
         out_states = self.token_dispatcher.combine(out_states)
         return out_states
+    
+    def forward_yield(self,
+                hidden_states: torch.Tensor,
+                topk_weights: torch.Tensor,
+                topk_ids: torch.LongTensor,
+                gate_up_weights: torch.Tensor,
+                gate_up_scale: torch.Tensor,
+                down_weights: torch.Tensor,
+                down_scale: torch.Tensor,
+                expert_list: List[int] = None,
+                tag: Any = None):
+        """forward_yield."""
+        topk_weights = _renormalize(topk_weights, self.renormalize)
+
+        _token_dispatcher = self.token_dispatcher
+        if tag is not None and tag[0] == "0":
+            _token_dispatcher = self.token_dispatcher
+        if tag is not None and tag[0] == "1":
+            _token_dispatcher = self.token_dispatcher_for2mb
+        recv_hidden_states, recv_topk_ids, recv_topk_weights, tokens_per_expert = (
+            yield from _token_dispatcher.dispatch_yield(
+                hidden_states,
+                topk_ids.to(torch.int64),
+                topk_weights.to(torch.float32),
+                self.num_experts,
+            )
+        )
+        out_states = self.experts.forward(recv_hidden_states, tokens_per_expert, gate_up_weights, gate_up_scale,
+                                 down_weights, down_scale)
+        out_states = yield from _token_dispatcher.combine_yield(out_states)
+        return out_states
 
 class TritonFusedMoEBlockedF8Builder(FusedMoEBlockedF8Builder):
     """triton fused moe blocked f8 builder."""

lmdeploy/pytorch/backends/cuda/token_dispatcher.py

@@ -181,7 +181,6 @@ def dispatch(
         num_max_dispatch_tokens_per_rank: int = 128,
     ) -> Tuple[torch.Tensor, torch.Tensor]:
         self.hidden_shape = hidden_states.shape
-        topk_idx = topk_idx.to(torch.int64)
         (
             hidden_states,
             topk_idx,
@@ -205,6 +204,39 @@ def dispatch(
             hidden_states = self.get_permuted_hidden_states_by_experts(hidden_states)
         return hidden_states, topk_idx, topk_weights, tokens_per_expert
 
+    def dispatch_yield(
+        self,
+        hidden_states: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+        num_experts: int,
+        previous_event=None,
+        num_max_dispatch_tokens_per_rank: int = 128,
+    ):
+        self.hidden_shape = hidden_states.shape
+        (
+            hidden_states,
+            topk_idx,
+            topk_weights,
+            num_recv_tokens_per_expert_list,
+            handle,
+            event,
+        ) = yield from self.dispatch_normal_yield(
+            hidden_states, topk_idx, topk_weights, num_experts, previous_event
+        )
+        self.tokens_per_expert = torch.tensor(
+            num_recv_tokens_per_expert_list,
+            device=hidden_states.device,
+            dtype=torch.int64,
+        )
+        tokens_per_expert = self.get_number_of_tokens_per_expert()
+        self.handle = handle
+        self.topk_idx = topk_idx
+        self.topk_weights = topk_weights
+        if hidden_states.shape[0] > 0:
+            hidden_states = self.get_permuted_hidden_states_by_experts(hidden_states)
+        return hidden_states, topk_idx, topk_weights, tokens_per_expert
+
     def dispatch_normal(
         self,
         x: torch.Tensor,
@@ -256,6 +288,62 @@ def dispatch_normal(
             event,
         )
 
+    def dispatch_normal_yield(
+        self,
+        x: torch.Tensor,
+        topk_idx: torch.Tensor,
+        topk_weights: torch.Tensor,
+        num_experts: int,
+        previous_event=None,
+        async_finish=True
+    ):
+        yield
+        previous_event = self.buffer_normal.capture() if async_finish else None
+        (
+            num_tokens_per_rank,
+            num_tokens_per_rdma_rank,
+            num_tokens_per_expert,
+            is_token_in_rank,
+            previous_event,
+        ) = self.buffer_normal.get_dispatch_layout(
+            topk_idx,
+            num_experts,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=previous_event is not None and async_finish,
+        )
+
+        (
+            recv_x,
+            recv_topk_idx,
+            recv_topk_weights,
+            num_recv_tokens_per_expert_list,
+            handle,
+            event,
+        ) = self.buffer_normal.dispatch(
+            x,
+            topk_idx=topk_idx,
+            topk_weights=topk_weights,
+            num_tokens_per_rank=num_tokens_per_rank,
+            num_tokens_per_rdma_rank=num_tokens_per_rdma_rank,
+            is_token_in_rank=is_token_in_rank,
+            num_tokens_per_expert=num_tokens_per_expert,
+            previous_event=previous_event,
+            async_finish=async_finish,
+            allocate_on_comm_stream=previous_event is not None and async_finish,
+        )
+
+        yield
+        if async_finish:
+            event.current_stream_wait()
+        return (
+            recv_x,
+            recv_topk_idx,
+            recv_topk_weights,
+            num_recv_tokens_per_expert_list,
+            handle,
+            event,
+        )
 
     def combine(
         self, hidden_states: torch.Tensor
@@ -268,6 +356,17 @@ def combine(
         self.handle = None
         return hidden_states.view(self.hidden_shape)
 
+    def combine_yield(
+        self, hidden_states: torch.Tensor
+    ):
+        if hidden_states.shape[0] > 0:
+            hidden_states = self.get_restored_hidden_states_by_experts(
+                hidden_states
+            )
+        hidden_states, event = yield from self.combine_normal_yield(hidden_states, self.handle)
+        self.handle = None
+        return hidden_states.view(self.hidden_shape)
+
     def combine_normal(self, x: torch.Tensor, handle: Tuple, previous_event=None):
         combined_x, _, event = self.buffer_normal.combine(
             x,
@@ -278,6 +377,22 @@ def combine_normal(self, x: torch.Tensor, handle: Tuple, previous_event=None):
         )
         return combined_x, event
 
+    def combine_normal_yield(self, x: torch.Tensor, handle: Tuple, previous_event=None, async_finish=True):
+        yield
+        previous_event = self.buffer_normal.capture() if async_finish else None
+        combined_x, _, event = self.buffer_normal.combine(
+            x,
+            handle,
+            async_finish=async_finish,
+            previous_event=previous_event,
+            allocate_on_comm_stream=previous_event is not None and async_finish,
+        )
+
+        yield
+        if async_finish:
+            event.current_stream_wait()
+        return combined_x, event
+
     def _indices_to_multihot(self, indices, probs):
         batch_size = indices.shape[0]
         multihot_routing_map = torch.zeros(

lmdeploy/pytorch/models/deepseek_v2.py

@@ -894,3 +894,5 @@ def __skip_nextn(name, nextn_keys):
                 else:
                     param = params_dict[name]
                     load_weight(param, loaded_weight)
+
+import lmdeploy.pytorch.models.utils.microbatch