Merge pull request #106 from raindaywhu/dev_whq_eplb2

collect moe load after dispatch

Author: wanghanqingLYT

vllm_ascend/eplb/adaptor/vllm_adaptor.py

@@ -84,42 +84,13 @@ def init_expert_param_per_layer(self):
                         for name in self.expert_weight_names]
                 )
 
-    def collect_topk_ids(self, dummy_run=False):
-        if dummy_run:
-            return 
-        self.all_topk_ids.append(self.model.get_all_topk_ids(self.num_moe_layers))
+    # def collect_topk_ids(self, dummy_run=False):
+    #     if dummy_run:
+    #         return
+        # self.all_topk_ids.append(self.model.get_all_topk_ids(self.num_moe_layers))
 
     def get_rank_expert_workload(self) -> torch.Tensor:
-        device = self.all_topk_ids[0][0].device
-        if not hasattr(self, "moe_load"):
-            self.moe_load = torch.zeros(
-                (self.num_moe_layers), self.global_expert_num,
-                dtype=torch.int64,
-                device=self.all_topk_ids[0][0].device,
-            )
-        else:
-            self.moe_load.zero_()   
-            # pass       
-        flat_list_per_layer = [[] for _ in range(self.num_moe_layers)]  
-
-        for period_data in self.all_topk_ids:     
-            for l in range(self.num_moe_layers):  
-                t = period_data[l]       
-                flat_list_per_layer[l].append(t.reshape(-1))  
-
-        index_2d = torch.nn.utils.rnn.pad_sequence(
-            [torch.cat(flat_list_per_layer[l]) for l in range(self.num_moe_layers)],
-            batch_first=True, padding_value=-1     
-        ).to(device)               
-
-        mask = index_2d != -1
-        index_2d = index_2d.masked_select(mask).reshape(self.num_moe_layers, -1)
-        src_2d   = torch.ones_like(index_2d, dtype=torch.int64)
-
-        self.moe_load.scatter_add_(dim=1, index=index_2d, src=src_2d)
-
-        if self.all_topk_ids:                     
-            self.all_topk_ids[:] = self.all_topk_ids[-1:]
+        self.moe_load = self.model.get_all_moe_loads()
         return self.moe_load
 
     def get_init_expert_map(self, num_moe_layers):
@@ -224,13 +195,13 @@ def determine_expert_map_all(self):
 
         for r in range(self.world_size):
             if r < self.world_size - 1:
-                start = r * local_num_experts 
-                end   = (r + 1) * local_num_experts 
-                local_count = local_num_experts 
+                start = r * local_num_experts
+                end   = (r + 1) * local_num_experts
+                local_count = local_num_experts
             else:
-                start = r * local_num_experts 
+                start = r * local_num_experts
                 end   = self.global_expert_num
-                local_count = self.global_expert_num - r * local_num_experts 
+                local_count = self.global_expert_num - r * local_num_experts
 
             local_ids = torch.arange(local_count, dtype=torch.int32)
             expert_map_all[:, r, start:end] = local_ids.unsqueeze(0).expand(self.num_moe_layers, -1)

vllm_ascend/eplb/core/worker/eplb_worker.py

@@ -61,7 +61,7 @@ def do_update(self):
             return
 
         #根据负载信息，获取更新后的专家表
-        load_info, old_placement = self.global2local(load_info, self.old_expert_maps, self.num_local_experts)
+        old_placement = self.global2local(self.old_expert_maps, self.num_local_experts)
         changed, priority, new_placement = self.calculate_rebalance_experts(load_info, old_placement)
 
         if not torch.is_tensor(new_placement):
@@ -276,18 +276,13 @@ def update_expert_map(self, expert_maps):
         self.shared_dict["expert_maps"] = expert_maps
 
     def global2local(self,
-        workload: torch.Tensor,
         placement: torch.Tensor,
         E_local: int
     ) -> tuple[torch.Tensor, torch.Tensor]:
 
         L, G, _ = placement.shape
         device = placement.device
 
-        wt_local = torch.full((L, G, E_local),
-                              fill_value=-1,
-                              dtype=workload.dtype,
-                              device=device)
         pt_local = torch.full((L, G, E_local),
                               fill_value=-1,
                               dtype=torch.long,
@@ -297,12 +292,10 @@ def global2local(self,
         l_idx, g_idx, k_idx = valid.nonzero(as_tuple=True)
 
         slot_idx = placement[l_idx, g_idx, k_idx]
-        values = workload[l_idx, g_idx, k_idx]
 
-        wt_local[l_idx, g_idx, slot_idx] = values
         pt_local[l_idx, g_idx, slot_idx] = k_idx
 
-        return wt_local, pt_local
+        return pt_local
 
 
     def local2global(self,

vllm_ascend/eplb/eplb_updator.py

@@ -80,8 +80,8 @@ def init_eplb(self, expert_map_path):
             shared_dict = self.shared_dict,
             planner_q = self.planner_block_queue,
             block_update_q = self.block_update_queue,
-            redundant_enable = self.redundant_enable, 
-            policy_type = 6,
+            redundant_enable = self.redundant_enable,
+            policy_type = 1,
             enable_d2d = True
         )
 
@@ -91,8 +91,8 @@ def init_eplb(self, expert_map_path):
 
     def get_update_iteration(self):
         self.cur_iterations = self.cur_iterations + 1
-        load_gather_iteration = self.cur_iterations % self.num_expert_load_gather == 0 if not self.gate_eplb else self.cur_iterations == self.num_iterations 
-        upate_iteration = self.cur_iterations % self.num_iterations == 0 if not self.gate_eplb else self.cur_iterations == self.num_iterations 
+        load_gather_iteration = self.cur_iterations % self.num_expert_load_gather == 0 if not self.gate_eplb else self.cur_iterations == self.num_iterations
+        upate_iteration = self.cur_iterations % self.num_iterations == 0 if not self.gate_eplb else self.cur_iterations == self.num_iterations
         return load_gather_iteration, upate_iteration
 
     def get_init_expert_map(self):
@@ -135,7 +135,6 @@ def forward_before(self):
 
 
     def forward_end(self,dummy_run=False):
-        self.adaptor.collect_topk_ids(dummy_run)
         if not self.update_in_flight:
             load_gather_iteration, update_iteration = self.get_update_iteration()
             if load_gather_iteration:
@@ -174,12 +173,12 @@ def compute_and_set_moe_load(self,dummy_run=False):
             moe_load = local_load.unsqueeze(1)
             self.shared_dict["moe_load"] = moe_load.cpu()
             logger.debug(f"[ModelRunner] Updated shared_dict['moe_load'] shape={moe_load.shape}")
+        self.adaptor.model.clear_all_moe_loads()
         return moe_load
 
     def warm_up_eplb(self):
 
         self.get_init_expert_map()
-        self.adaptor.collect_topk_ids(dummy_run=False)
         self.compute_and_set_moe_load()
 
         src_tensor = torch.empty((1,), device=self.device)
@@ -240,7 +239,6 @@ def get_expert_load(self) -> tuple:
         num_local_experts = expert_maps.max() + 1
         return  moe_load, expert_maps, num_local_experts
 
-
     def update_expert_load_statistical_period(self, num_expert_load_gather: int, num_iterations: int):
         logger.info(f" start update {self.num_expert_load_gather=}, {self.num_iterations}...")
         self.num_expert_load_gather = num_expert_load_gather

vllm_ascend/models/deepseek_v2.py

@@ -727,6 +727,9 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         quant_config = vllm_config.quant_config
         self.config = config
         self.quant_config = quant_config
+        self.num_dense_layers = self.config.first_k_dense_replace
+        self.num_moe_layers = self.config.num_hidden_layers - self.num_dense_layers
+
         self.model = CustomDeepseekV2Model(vllm_config=vllm_config,
                                            prefix=maybe_prefix(
                                                prefix, "model"))
@@ -755,29 +758,31 @@ def forward(
                                    inputs_embeds)
         return hidden_states
 
-    def get_expert_map(self,layer_id):
+    def get_expert_map(self, layer_id):
         return self.model.layers[layer_id].mlp.experts.get_map()
 
-    def get_log2phy_map(self,layer_id):
+    def get_log2phy_map(self, layer_id):
         return self.model.layers[layer_id].mlp.experts.get_log2phy_map()
 
-    def get_all_expert_map(self,num_moe_layers):
+    def get_all_expert_map(self, num_moe_layers):
         all_loads = []
         for layer_id in range(num_moe_layers):
             load_tensor = self.get_expert_map(3+layer_id)  # (num_experts_per_layer,)
             all_loads.append(load_tensor)
 
         return torch.stack(all_loads, dim=0)
 
-    def get_topk_ids(self,layer_id):
-        return self.model.layers[layer_id+3].mlp.experts.topk_ids
+    def get_all_moe_loads(self):
+        all_moe_loads = torch.stack(
+            [self.model.layers[layer_id + self.num_dense_layers].mlp.experts.moe_load \
+                for layer_id in range(self.num_moe_layers)],
+            dim=0
+        )
+        return all_moe_loads
 
-    def get_all_topk_ids(self,num_moe_layers):
-        all_topk_id = []
-        for layer_id in range(num_moe_layers):
-            load_tensor = self.get_topk_ids(layer_id)  
-            all_topk_id.append(load_tensor)
-        return all_topk_id
+    def clear_all_moe_loads(self):
+        for layer_id in range(self.num_moe_layers):
+            self.model.layers[layer_id + self.num_dense_layers].mlp.experts.clear_moe_load()
 
 class CustomDeepseekV3ForCausalLM(CustomDeepseekV2ForCausalLM):
     pass

vllm_ascend/ops/fused_moe.py

@@ -1012,8 +1012,6 @@ def __init__(
 
         AscendFusedMoE.moe_counter += 1
         self.moe_instance_id = AscendFusedMoE.moe_counter
-        self.moe_load = None
-        self.topk_ids =  None
 
         if params_dtype is None:
             params_dtype = torch.get_default_dtype()
@@ -1103,6 +1101,10 @@ def __init__(
         local_num_experts = torch.sum(self.expert_map != -1) \
             if self.expert_map is not None else num_experts
 
+        self.moe_load = None
+        if self.dynamic_eplb:
+            self.moe_load = torch.zeros(local_num_experts, dtype=torch.int64)
+
         moe_quant_params = {
             "num_experts": local_num_experts,
             "hidden_size": hidden_size,
@@ -1176,7 +1178,7 @@ def forward(self,
             router_logits = get_dp_group().all_gather(router_logits, 0)
 
         # Matrix multiply.
-        e_hidden_states, self.topk_ids = self.quant_method.apply(
+        e_hidden_states, expert_token_num, group_list_type = self.quant_method.apply(
             layer=self,
             x=hidden_states,
             router_logits=router_logits,
@@ -1198,6 +1200,10 @@ def forward(self,
             and self.enable_multistream_moe and not is_prefill else None,
         )
 
+        if self.dynamic_eplb:
+            self.moe_load += expert_token_num if group_list_type else \
+                torch.cat([expert_token_num[:1], expert_token_num[1:] - expert_token_num[:-1]])
+
         if shared_experts:
             if isinstance(e_hidden_states, tuple):
                 e_hidden_states, shared_hidden_states = e_hidden_states
@@ -1267,3 +1273,7 @@ def get_map(self):
     def get_log2phy_map(self):
         return self.log2phy
 
+    def clear_moe_load(self):
+        self.moe_load.zero_()
+
+

vllm_ascend/quantization/w8a8_dynamic.py

@@ -208,13 +208,14 @@ def fused_experts_with_mc2(
 
     hidden_states = torch_npu.npu_moe_distribute_combine(**kwargs_mc2)
 
+    group_list_type = 1
     if shared_experts is None:
-        return hidden_states
+        return hidden_states, expert_token_nums, group_list_type
     else:
         with npu_stream_switch("moe_secondary", 0):
             npu_wait_tensor(shared_act[0], down_out_list)
             shared_output, _ = shared_experts.down_proj(shared_act)
-        return hidden_states, shared_output
+        return hidden_states, shared_output, expert_token_nums, group_list_type
 
 
 # currently expert parallelism implemented with all2all
@@ -343,7 +344,7 @@ def fused_experts_with_all2all(
         )
     if len(original_shape) == 3:
         final_hidden_states = final_hidden_states.view(original_shape)
-    return final_hidden_states
+    return final_hidden_states, expert_tokens, group_list_type
 
 
 def fused_experts(hidden_states: torch.Tensor,
@@ -457,7 +458,7 @@ def fused_experts(hidden_states: torch.Tensor,
 
     if len(original_shape) == 3:
         final_hidden_states = final_hidden_states.view(original_shape)
-    return final_hidden_states
+    return final_hidden_states, expert_tokens, group_list_type
 
 
 class AscendW8A8DynamicLinearMethod:
@@ -677,7 +678,7 @@ def apply(
                 log2phy=log2phy,
                 global_redundant_expert_num=global_redundant_expert_num,
                 shared_experts=shared_experts,
-                **kwargs), topk_ids
+                **kwargs)
         elif fused_moe_state == FusedMoEState.AllGather:
             return fused_experts(hidden_states=x,
                                  w1=layer.w13_weight,
@@ -687,7 +688,7 @@ def apply(
                                  topk_weights=topk_weights,
                                  topk_ids=topk_ids,
                                  top_k=top_k,
-                                 expert_map=expert_map), topk_ids
+                                 expert_map=expert_map)
         else:
             # The current implementation of deepseek moe splits hidden_states
             # according to tp_size before they are feed into fused_moe module.
@@ -706,7 +707,7 @@ def apply(
                 ep_group=self.ep_group,
                 log2phy=log2phy,
                 global_redundant_expert_num=global_redundant_expert_num,
-            ), topk_ids
+            )
 
     def process_weights_after_loading(self, layer):
         if self.transpose_weight:

vllm_ascend/worker/model_runner_v1.py

@@ -366,7 +366,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.dp_size = vllm_config.parallel_config.data_parallel_size
         self.dp_rank = vllm_config.parallel_config.data_parallel_rank
 
-        #EPLB 
+        #EPLB
         self.dynamic_eplb = ascend_config.dynamic_eplb
         if self.dynamic_eplb == True:
             self.eplb_adaptor = None
@@ -1240,7 +1240,7 @@ def execute_model(
 
             if self.dynamic_eplb:
                 self.eplb_updator.forward_before()
-              
+
             (attn_metadata, hidden_states, spec_decode_metadata, positions,
              num_scheduled_tokens,
              sample_indices) = (self._process_reqs(scheduler_output,
@@ -1544,6 +1544,8 @@ def _dummy_run(
                         intermediate_tensors=intermediate_tensors,
                         inputs_embeds=inputs_embeds)
 
+                if is_profile_run and self.dynamic_eplb:
+                    self.model.clear_all_moe_loads()
                 if not is_compile and not is_profile_run and self.dynamic_eplb:
                     dummy_run = True
                     self.eplb_updator.forward_end(dummy_run)