Merge pull request #102 from raindaywhu/br_main_into_eplb_wjh

fix bug in moe load & add expert load to josn

Author: qmkakaxi

vllm_ascend/eplb/adaptor/vllm_adaptor.py

@@ -90,8 +90,16 @@ def collect_topk_ids(self, dummy_run=False):
         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:     
@@ -108,12 +116,11 @@ def get_rank_expert_workload(self) -> torch.Tensor:
         index_2d = index_2d.masked_select(mask).reshape(self.num_moe_layers, -1)
         src_2d   = torch.ones_like(index_2d, dtype=torch.int64)
 
-        moe_load = torch.zeros((self.num_moe_layers),  self.global_expert_num,
-                            dtype=torch.int64, device=device)
-        moe_load.scatter_add_(dim=1, index=index_2d, src=src_2d)
+        self.moe_load.scatter_add_(dim=1, index=index_2d, src=src_2d)
 
-        self.all_topk_ids = []
-        return moe_load
+        if self.all_topk_ids:                     
+            self.all_topk_ids[:] = self.all_topk_ids[-1:]
+        return self.moe_load
 
     def get_init_expert_map(self, num_moe_layers):
         expert_map = self.model.get_all_expert_map(num_moe_layers)
@@ -135,32 +142,6 @@ def get_init_expert_map(self, num_moe_layers):
 
         return all_expert_maps
 
-    def local2global(self,
-        placement_local: torch.Tensor
-    ) -> torch.Tensor:
-
-        L, G, E_local = placement_local.shape
-        device = placement_local.device
-
-        max_id = torch.max(placement_local)
-        E_global = (max_id + 1).item() if max_id >= 0 else 0
-
-        if E_global == 0:
-            return torch.empty((L, G, 0), dtype=torch.long, device=device)
-
-        placement_global = torch.full((L, G, E_global),
-                                      fill_value=-1,
-                                      dtype=torch.long,
-                                      device=device)
-
-        valid = placement_local >= 0
-        l_idx, g_idx, slot_idx = valid.nonzero(as_tuple=True)
-        gid_idx = placement_local[l_idx, g_idx, slot_idx]
-
-        placement_global[l_idx, g_idx, gid_idx] = slot_idx
-
-        return placement_global
-
     def get_init_expert_map_from_file(self, num_moe_layers, expert_map_path):
 
         try:

vllm_ascend/eplb/core/worker/eplb_worker.py

@@ -62,6 +62,7 @@ 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):

vllm_ascend/eplb/eplb_updator.py

@@ -14,14 +14,17 @@
 # limitations under the License.
 # This file is a part of the vllm-ascend project.
 #
+
 import torch
+from typing import Dict, List
 import torch.distributed as dist
 import vllm.envs as envs
 from multiprocessing import Queue, Manager
 
 from vllm.logger import logger
 from vllm_ascend.eplb.core.worker.eplb_worker import EplbProcess
 from vllm_ascend.eplb.core.loader.device_transfer_loader import D2DExpertWeightLoader
+from vllm_ascend.eplb.tool.eplb_utils import ExpertMapUtils
 
 class EplbUpdator:
 
@@ -32,6 +35,7 @@ def set_adaptor(self, adaptor):
         self.adaptor = adaptor
         self.eplb_loader = D2DExpertWeightLoader(eplb_adaptor=self.adaptor)
         self.num_moe_layers = self.adaptor.num_moe_layers
+        self.global_expert_num = self.adaptor.global_expert_num
 
     def init_eplb(self, expert_map_path):
         self.num_expert_load_gather = 10
@@ -69,7 +73,9 @@ def init_eplb(self, expert_map_path):
             # 热度负载信息 [num_layers, world_size, num_experts]
             "moe_load": None,
             # 所有的专家表[num_layers, world_size, num_experts]
-            "expert_maps": None
+            "expert_maps": None,
+            # 热度负载信息 [num_layers, world_size, local_num_experts]
+            "load_info": None,
         })
 
         self.eplb = EplbProcess(
@@ -125,11 +131,11 @@ def forward_before(self):
                 self.weight_update_counter = 0
                 self.update_in_flight = False
                 self.update_info_all = []
-
         # set asynchronous stream for d2d expert weight update
         self.reqs = []
         self.eplb_loader.asyn_expert_weight_transfer(self.reqs)
 
+
     def forward_end(self,dummy_run=False):
         self.adaptor.collect_topk_ids(dummy_run)
         if not self.update_in_flight:
@@ -149,6 +155,7 @@ def forward_end(self,dummy_run=False):
 
     def compute_and_set_moe_load(self,dummy_run=False):
         local_load = self.adaptor.get_rank_expert_workload()
+
         self._gather_buffer = None
         if dist.is_initialized():
             self.world_size = dist.get_world_size()
@@ -229,28 +236,31 @@ def unpack_update_batch(self, packed_update_info):
         return recovered
 
     def get_expert_load(self) -> str:
-
-        # todo wjh 给到返回值
-        # return self.shared_dict['moe_load']
-        # mock json_str
-        experts_load = ('{\"expert_load\":['
-                        '{\"ip\":\"141.xxx.xxx.181\",'
-                        '\"node_0\":'
-                        '{\"card_0\":'
-                        '[{\"layer_4\":{\"expert_0\":3,\"expert_2\":1}},{\"layer_5\":{\"expert_0\":3,\"expert_2\":1}}],'
-                        '\"card_1\":[{\"layer_4\":{\"expert_1\":3,\"expert_3\":1},\"layer_5\":{\"expert_0\":3,\"'
-                        'expert_2\":1}}]}},{\"ip\":\"141.xxx.xxx.177\",\"node_0\":{\"card_0\":[{\"layer_4\":'
-                        '{\"expert_0\":3,\"expert_2\":1}},{\"layer_5\":{\"expert_0\":3,\"expert_2\":1}}],'
-                        '\"card_1\":[{\"layer_4\":{\"expert_1\":3,\"expert_3\":1}}]}}]}')
-        return experts_load
+        
+        load_info = self.shared_dict["load_info"]  # Tensor [L, W, local_experts_num]  
+        L, W, _ = load_info.shape
+
+        expert_load: Dict[str, List[dict]] = {}
+        for c in range(W):
+            layers: List[dict] = []
+            for l in range(L):
+                counts_1d = load_info[l, c]         
+        
+                layer_val = {
+                    f"expert_{e}": int(v)            
+                    for e, v in enumerate(counts_1d.tolist())
+                }
+                layers.append({f"layer_{l}": layer_val})
+            expert_load[f"card_{c}"] = layers
+
+        return {"expert_load": expert_load}
 
     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
         self.num_iterations = num_iterations
         logger.info(f" update {self.num_expert_load_gather=}, {self.num_iterations} success...")
 
-
     def shutdown(self):
         """
         Clean up the EPLB process.