Merge pull request #105 from raindaywhu/br_main_into_eplb_wjh

fix get_expert_load

Author: qmkakaxi

vllm_ascend/eplb/core/worker/eplb_worker.py

@@ -62,14 +62,13 @@ def do_update(self):
 
         #根据负载信息，获取更新后的专家表
         load_info, old_placement = self.global2local(load_info, self.old_expert_maps, self.num_local_experts)
-        self.shared_dict["load_info"] = load_info
         changed, priority, new_placement = self.calculate_rebalance_experts(load_info, old_placement)
 
         if not torch.is_tensor(new_placement):
             new_placement = torch.tensor(new_placement)
         self.check_expert_placement(old_placement, new_placement)
         new_expert_maps = self.local2global(new_placement)
-
+        self.update_expert_map(new_expert_maps)
         logger.debug(f"[EPLB Process  new_map differs, performing D2D")
 
         update_info = self.compose_expert_update_info_bipartite(new_expert_maps, self.old_expert_maps)\

vllm_ascend/eplb/eplb_updator.py

@@ -74,8 +74,6 @@ def init_eplb(self, expert_map_path):
             "moe_load": None,
             # 所有的专家表[num_layers, world_size, num_experts]
             "expert_maps": None,
-            # 热度负载信息 [num_layers, world_size, local_num_experts]
-            "load_info": None,
         })
 
         self.eplb = EplbProcess(
@@ -235,9 +233,12 @@ def unpack_update_batch(self, packed_update_info):
         ]
         return recovered
 
-    def get_expert_load(self) -> str:
+    def get_expert_load(self):
+        expert_maps = self.shared_dict["expert_maps"]
+        moe_load = self.shared_dict["moe_load"]  # Tensor [L, W, global_experts_num]  
+        num_local_experts = expert_maps.max() + 1
+        load_info, _ = ExpertMapUtils.global2local_load(moe_load, expert_maps, num_local_experts)
 
-        load_info = self.shared_dict["load_info"]  # Tensor [L, W, local_experts_num]  
         L, W, _ = load_info.shape
 
         expert_load: Dict[str, List[dict]] = {}

vllm_ascend/eplb/tool/eplb_utils.py

@@ -83,3 +83,33 @@ def global2local(cls,
         pt_local[g_idx, slot_idx] = k_idx
 
         return pt_local
+
+    @classmethod
+    def global2local_load(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,
+                              device=device)
+
+        valid = placement >= 0
+        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