add fix warm up & change init expert map from file

Author: qmkakaxi

vllm_ascend/eplb/adaptor/vllm_adaptor.py

@@ -121,6 +121,27 @@ def get_init_expert_map(self, num_moe_layers):
 
         return all_expert_maps
 
+    def get_init_expert_map_from_file(self, num_moe_layers, expert_map_path):
+        expert_map_tensor, layers_num, ranks_num = self._expert_file_to_tensor(expert_map_path)
+        for layer_idx in range(num_moe_layers):
+            self.expert_map_per_layer_cpu[layer_idx] = \
+                expert_map_tensor[layer_idx][self.rank_id]
+    
+    def _expert_file_to_tensor(self, expert_map_path: str):
+        with open(expert_map_path, "r") as f:
+            data = json.load(f)
+        layers_num = data["moe_layer_count"]
+        gpus_num = data["layer_list"][0]["device_count"]
+
+        tensor_data = []
+        for layer in data["layer_list"]:
+            device_data = []
+            for device in layer["device_list"]:
+                device_data.append(device["device_expert"])
+            tensor_data.append(device_data)
+        expert_map_tensor = torch.tensor(tensor_data, dtype=torch.int32)
+        return expert_map_tensor, layers_num, gpus_num
+
     def do_update_expert_map(self, layer_id, updated_expert_map):
         self.expert_map_per_layer[layer_id].copy_(updated_expert_map)
         self.expert_map_per_layer_cpu[layer_id].copy_(updated_expert_map)

vllm_ascend/eplb/eplb_updator.py

@@ -24,18 +24,19 @@
 
 class EplbUpdator:
 
-    def __init__(self, redundant_enable):
-        self.init_eplb(redundant_enable)
+    def __init__(self, expert_map_path):
+        self.init_eplb(expert_map_path)
 
     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
 
-    def init_eplb(self, redundant_enable):
+    def init_eplb(self, expert_map_path):
 
-        self.redundant_enable = redundant_enable 
+        self.redundant_enable = (expert_map_path != None)
         self.num_iterations: torch.int64 = 130
+        self.expert_map_path = expert_map_path
 
         self.weight_update_counter = 0
         self.expert_map_initialized = False
@@ -82,7 +83,8 @@ def get_update_iteration(self):
     def get_init_expert_map(self):
         try:
             if not self.expert_map_initialized:
-                self.shared_dict["expert_maps"] = self.adaptor.get_init_expert_map(self.num_moe_layers)
+                # self.shared_dict["expert_maps"] = self.adaptor.get_init_expert_map(self.num_moe_layers)
+                self.shared_dict["expert_maps"] = self.adaptor.get_init_expert_map_from_file(self.num_moe_layers, self.expert_map_path)
                 self.expert_map_initialized = True
         except Exception as e:
             logger.warning(f"[ModelRunner] Failed to wake EPLB process: {e}", exc_info=True)

vllm_ascend/worker/model_runner_v1.py

@@ -371,7 +371,7 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device):
         self.dynamic_eplb = ascend_config.dynamic_eplb
         if self.dynamic_eplb == True:
             self.eplb_adaptor = None
-            self.eplb_updator = EplbUpdator(ascend_config.expert_map_path != None)
+            self.eplb_updator = EplbUpdator(ascend_config.expert_map_path)
 
     def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
         """Update the cached states and the persistent batch with the scheduler
@@ -1508,12 +1508,6 @@ def _dummy_run(
                         intermediate_tensors=intermediate_tensors,
                         inputs_embeds=inputs_embeds)
 
-                #EPLB
-                if self.dynamic_eplb == True:
-                    self.eplb_adaptor = VllmEplbAdaptor(model=self.model)
-                    self.eplb_updator.set_adaptor(self.eplb_adaptor)
-                    self.eplb_updator.warm_up_eplb()
-
                 return hidden_states
 
     def profile_run(self) -> None:
@@ -1555,6 +1549,13 @@ def profile_run(self) -> None:
         self.encoder_cache.clear()
         gc.collect()
 
+    def eplb_warmup(self):
+        #EPLBMore actions
+        if self.dynamic_eplb == True:
+            self.eplb_adaptor = VllmEplbAdaptor(model=self.model)
+            self.eplb_updator.set_adaptor(self.eplb_adaptor)
+            self.eplb_updator.warm_up_eplb()
+
     def load_model(self) -> None:
         logger.info("Starting to load model %s...", self.model_config.model)
 

vllm_ascend/worker/worker_v1.py

@@ -203,6 +203,7 @@ def compile_or_warm_up_model(self) -> None:
         for size in sorted(warmup_sizes, reverse=True):
             logger.info("Compile and warming up model for size %d", size)
             self.model_runner._dummy_run(size)
+            self.model_runner.eplb_warmup()
         if not self.model_config.enforce_eager:
             self.model_runner.capture_model()
         # Reset the seed to ensure that the random state is not affected by