Merge pull request #99 from raindaywhu/lt_expert_load

expert load collecting

Author: raindaywhu

vllm_ascend/eplb/eplb_updator.py

@@ -140,6 +140,7 @@ def forward_end(self,dummy_run=False):
             load_gather_iteration, update_iteration = self.get_update_iteration()
             if load_gather_iteration:
                 moe_load = self.compute_and_set_moe_load()
+                self.get_expert_load()
             if update_iteration:
                 self.wakeup_eplb_worker()
                 self.update_in_flight = True
@@ -233,28 +234,12 @@ def unpack_update_batch(self, packed_update_info):
         ]
         return recovered
 
-    def get_expert_load(self):
+    def get_expert_load(self) -> tuple:
         expert_maps = self.shared_dict["expert_maps"]
-        moe_load = self.shared_dict["moe_load"]  # Tensor [L, W, global_experts_num]  
+        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)
-        
-        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}
+        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}...")

vllm_ascend/eplb/tool/eplb_utils.py

@@ -90,7 +90,6 @@ def global2local_load(self,
         placement: torch.Tensor,
         E_local: int
     ) -> tuple[torch.Tensor, torch.Tensor]:
-
         L, G, _ = placement.shape
         device = placement.device
 

vllm_ascend/worker/model_runner_v1.py

@@ -1588,7 +1588,7 @@ def profile_run(self) -> None:
         self.encoder_cache.clear()
         gc.collect()
 
-    def do_get_expert_load(self) -> str:
+    def do_get_expert_load(self) ->  tuple:
         return self.eplb_updator.get_expert_load()
 
     def do_update_expert_load_statistical_period(self, num_expert_load_gather: int, num_iterations: int):

vllm_ascend/worker/worker_v1.py

@@ -209,11 +209,8 @@ def compile_or_warm_up_model(self) -> None:
         # the model initialization and profiling.
         set_random_seed(self.model_config.seed)
 
-    def get_expert_load(self) -> str:
-        """ todo 一共几个worker"""
-        moe_load = self.model_runner.do_get_expert_load()
-        return moe_load
-
+    def get_expert_load(self) -> tuple:
+        return self.model_runner.do_get_expert_load()
     def update_expert_load_statistical_period(self, num_expert_load_gather: int, num_iterations: int):
         self.model_runner.do_update_expert_load_statistical_period(num_expert_load_gather, num_iterations)