merge from remote default branch and fix conflict

Author: lt

vllm_ascend/eplb/core/policy/dynamic_ep_v2.py

@@ -383,7 +383,30 @@ def non_redundant_expert_information(origin_deployment, updated_weights, num_rad
         return device_assignments, device_weights, device_loads, device_counts
 
     @staticmethod
-    def distribute_redun_experts(device_assignments, device_weights, device_loads, device_counts, redundant_expert_list,
+    def recomputing_weight(layer_workloads, device_assignments, device_weights, device_loads):
+        # 统计专家出现次数
+        num_all_experts = [0] * len(layer_workloads)
+        num_devices = len(device_assignments)
+        for device_id in range(num_devices):
+            num_expert_per_npu = len(device_assignments[device_id])
+            for idx in range(num_expert_per_npu):
+                num_all_experts[idx] += device_assignments[device_id][idx]
+
+        for device_id in range(num_devices):
+            num_expert_per_npu = len(device_weights[device_id])
+            total_weight = 0.0
+            for idx in range(num_expert_per_npu):
+                expert_id = device_assignments[device_id][idx]
+                if num_all_experts[expert_id] == 0:
+                    print("Error: Division by zero")
+                device_weights[device_id][idx] = layer_workloads[expert_id] / num_all_experts[expert_id]
+                total_weight += device_weights[device_id][idx]
+            device_loads[device_id] = total_weight
+
+        return device_weights, device_loads
+
+    @staticmethod
+    def distribute_redun_experts(self, layer_workloads, device_assignments, device_weights, device_loads, device_counts, redundant_expert_list,
                                  items_per_device, expert_form_device, num_experts):
 
         num_devices = len(device_assignments)
@@ -411,18 +434,16 @@ def distribute_redun_experts(device_assignments, device_weights, device_loads, d
                 communication_box_index = expert_form_device[expert_id]
                 com_between_devices[candidate][communication_box_index] = expert_id
         # 极端情况下存在冗余专家没装箱 导致箱子有空位 随机填入专家 待优化
+        flag = False
         for dev_id in range(num_devices):
             # 检查容量限制
             if device_counts[dev_id] < items_per_device:
                 # 遍历合适的专家
                 for expert_id in range(num_experts):
                     if expert_id not in device_assignments[dev_id]:
-                        # 找到对应权重
-                        weight = 0
-                        for i in range(num_devices):
-                            for j in range(len(device_assignments[i])):
-                                if expert_id == device_assignments[i][j]:
-                                    weight = device_weights[i][j]
+                        flag = True
+                        # 随机初始化一个权重
+                        weight = 0.0
                         # 和该专家相关的卡权重发生变化 待修改
                         device_assignments[dev_id].insert(0, expert_id)
                         device_weights[dev_id].insert(0, weight)
@@ -432,12 +453,14 @@ def distribute_redun_experts(device_assignments, device_weights, device_loads, d
                         communication_box_index = expert_form_device[expert_id]
                         com_between_devices[dev_id][communication_box_index] = expert_id
                         break
-        #todo 重新生成权重
+
+        if flag:
+            device_weights, device_loads = self.recomputing_weight(layer_workloads, device_assignments, device_weights, device_loads)
 
         return device_assignments, device_weights, device_loads, device_counts, com_between_devices
 
     @staticmethod
-    def redundancy_again(self, origin_weights, num_redundant_experts, origin_deployment, expert_form_device, num_node,
+    def redundancy_again(self, layer_workloads, origin_weights, num_redundant_experts, origin_deployment, expert_form_device, num_node,
                          is_node_redundant):
 
         # 每张卡上专家数量
@@ -461,6 +484,8 @@ def redundancy_again(self, origin_weights, num_redundant_experts, origin_deploym
 
         # 新计算的冗余专家进行分配
         device_assignments, device_weights, device_loads, device_counts, com_between_devices = self.distribute_redun_experts(
+            self,
+            layer_workloads,
             device_assignments,
             device_weights,
             device_loads,
@@ -554,6 +579,7 @@ def redundant_expert_deployment(self, layer_workloads, original_deployment, expe
 
                 cur_device_assignments, cur_device_weights, cur_device_loads, cur_device_counts, cur_com_between_devices = self.redundancy_again(
                     self,
+                    layer_workloads,
                     cur_node_weights,
                     per_node_redun_expert_num,
                     cur_original_deployment,
@@ -569,6 +595,7 @@ def redundant_expert_deployment(self, layer_workloads, original_deployment, expe
         else:
             device_assignments, device_weights, device_loads, device_counts, com_between_devices = self.redundancy_again(
                 self,
+                layer_workloads,
                 weights,
                 redundancy_expert_num,
                 original_deployment,
@@ -583,7 +610,7 @@ def redundant_expert_deployment(self, layer_workloads, original_deployment, expe
 
     @staticmethod
     def two_device_exchange_experts(cur_device_result, exchange_device_result, cur_exchanged_expert_id,
-                                    next_exchanged_expert_id, ave_workload, increment, num_redundancy_expert):
+                                    next_exchanged_expert_id, ave_workload, increment, num_redundancy_expert, cur_org_placement, next_org_placement):
 
         cur_device_weight = cur_device_result['expert_weights']
         next_device_weight = exchange_device_result['expert_weights']
@@ -609,7 +636,8 @@ def two_device_exchange_experts(cur_device_result, exchange_device_result, cur_e
                     continue
                     # 交换专家限制卡内专家不同
                 change_flag = True
-                if cur_device_expert_id[index] in next_device_expert_id or next_device_expert_id[next_index] in cur_device_expert_id:
+                if ((cur_device_expert_id[index] in next_device_expert_id or next_device_expert_id[next_index] in cur_device_expert_id) or
+                        (cur_org_placement[0] == next_device_expert_id[next_index] or next_org_placement[0] == cur_device_expert_id[index])):
                     change_flag = False
                 # 选择的专家不能是参与过交换的
                 if (cur_device_expert_id[index] not in cur_exchanged_expert_id) and (
@@ -627,8 +655,7 @@ def two_device_exchange_experts(cur_device_result, exchange_device_result, cur_e
 
     @staticmethod
     def expert_exchange_between_devices(self, ave_workload, increment, cur_layer_result, com_between_devices, num_redundancy_expert,
-                                        node_idx=0,
-                                        per_node_device_num=0, is_node_redundant=False):
+                                        org_placement_table, node_idx=0, per_node_device_num=0, is_node_redundant=False):
 
         if is_node_redundant:
             # 拿出当前节点内设备的信息
@@ -677,7 +704,9 @@ def expert_exchange_between_devices(self, ave_workload, increment, cur_layer_res
                         next_exchanged_expert_id,
                         ave_workload,
                         increment,
-                        num_redundancy_expert)
+                        num_redundancy_expert,
+                        org_placement_table[max_weight_device_id],
+                        org_placement_table[min_weight_device_id])
 
                     # 有符合条件的专家进行交换
                     if cur_exchange_index != -1:
@@ -700,7 +729,7 @@ def expert_exchange_between_devices(self, ave_workload, increment, cur_layer_res
 
     @staticmethod
     def exchange_experts(self, layer_result, layer_com_between_devices, num_nodes, device_num, is_node_redundant,
-                         ave_workload, increment, num_redundancy_expert):
+                         ave_workload, increment, num_redundancy_expert, org_placement_table):
 
         global_deployment = []
 
@@ -709,9 +738,9 @@ def exchange_experts(self, layer_result, layer_com_between_devices, num_nodes, d
             for node_idx in range(num_nodes):
                 self.expert_exchange_between_devices(self, ave_workload, increment, layer_result,
                                                      layer_com_between_devices, num_redundancy_expert,
-                                                     node_idx, per_node_device_num, is_node_redundant)
+                                                     org_placement_table, node_idx, per_node_device_num, is_node_redundant)
         else:
-            self.expert_exchange_between_devices(self, ave_workload, increment, layer_result, layer_com_between_devices, num_redundancy_expert)
+            self.expert_exchange_between_devices(self, ave_workload, increment, layer_result, layer_com_between_devices, num_redundancy_expert, org_placement_table)
 
         max_workload = 0
         for box in layer_result:
@@ -734,14 +763,15 @@ def count_elements(self, lst):
         return count
 
     def rebalance_experts(self, current_expert_table, expert_workload):
-
+        # 输入：当前专家部署信息和对应的负载信息，形状为layer_num, num_npus, experts_per_npu
         info = DynamicTable()
-        info.workload_table = np.array(expert_workload)
-        info.placement_table = np.array(current_expert_table)
+        info.workload_table = expert_workload.numpy()
+        info.placement_table = current_expert_table.numpy()
         layer_num, num_npus, experts_per_npu = info.workload_table.shape
         expert_ids, counts = np.unique(info.placement_table[0], return_counts=True)
         num_redundancy_expert = self.get_redundant_num(num_npus, counts)
         num_original_expert = len(expert_ids)
+        # 负载信息转化为 58 * 256
         layer_workloads = self.add_redundant(info.placement_table, info.workload_table, num_original_expert)
         max_heat_per_layer_before = self.calculate_max_heat_per_layer(info.workload_table, layer_num)
         npu_heat_all_origin = sum(max_heat_per_layer_before)
@@ -764,50 +794,31 @@ def rebalance_experts(self, current_expert_table, expert_workload):
         # 每个卡部署的专家数量 一个冗余专家
         global_deployment = [[[] for _ in range(num_npus)] for _ in range(layer_num)]
         # 统计更换数据集后的初始58层不均衡度
-        layer_initial_imbalance = self.calculate_initial_imbalance(current_expert_table, layer_workloads)
+        layer_initial_imbalance = self.calculate_initial_imbalance(info.placement_table, layer_workloads)
         # 遍历获得每一层的放置策略，考虑计算均衡
         max_heat_per_layer_after = np.zeros([layer_num])
         sum_num = 0
         for layer in range(layer_num):
             # 不均衡度小于特定阈值不调整
             if layer_initial_imbalance[layer] < 1.1:
-                global_deployment[layer] = current_expert_table[layer]
+                global_deployment[layer] = info.placement_table[layer]
                 continue
 
             ave_workload = np.sum(layer_workloads[layer]) / num_npus
-            for device_id, device in enumerate(current_expert_table[layer]):
+            for device_id, device in enumerate(info.placement_table[layer]):
                 for index, expert_id in enumerate(device):
                     if index != 0:
                         expert_from_device[layer][expert_id] = device_id
 
             # 调整冗余专家
             result, max_workload, com_between_devices = self.redundant_expert_deployment(self, layer_workloads[layer],
-                                                                                         current_expert_table[layer],
+                                                                                         info.placement_table[layer],
                                                                                          expert_from_device[layer],
                                                                                          num_node, False)
             # 交换专家
             global_deployment[layer], new_max_workload = self.exchange_experts(self, result, com_between_devices,
                                                                                num_node, num_npus, False, ave_workload,
-                                                                               0.05, num_redundancy_expert)
-
-            # To guarantee there is no expert movement inside a NPU
-            start_physical_idx = 1 if num_redundancy_expert else 0
-            for rank in range(num_npus):
-                physical_expert = start_physical_idx
-                while physical_expert in range(start_physical_idx, experts_per_npu):
-                    # skip the expert which is moved into this rank
-                    if global_deployment[layer][rank][physical_expert] not in current_expert_table[layer, rank, :]:
-                        physical_expert += 1
-                        continue
-
-                    if global_deployment[layer][rank][physical_expert] != current_expert_table[layer][rank][physical_expert]:
-                        right_idx = np.where(current_expert_table[layer][rank] == global_deployment[layer][rank][physical_expert])[0][0]
-                        # exchange expert with the expert on the right physical index
-                        tempt = global_deployment[layer][rank][right_idx]
-                        global_deployment[layer][rank][right_idx] = global_deployment[layer][rank][physical_expert]
-                        global_deployment[layer][rank][physical_expert] = tempt
-                    else:
-                        physical_expert += 1
+                                                                               0.05, num_redundancy_expert, info.placement_table[layer])
 
             for device_id in range(num_npus):
                 com_between_devices[device_id] = {int(key): int(value) for key, value in
@@ -828,8 +839,4 @@ def rebalance_experts(self, current_expert_table, expert_workload):
         if npu_heat_all_after < 0.95 * npu_heat_all_origin:
             change = 1
 
-        return change, per_layer_priority, np.array(global_deployment).tolist()
-
-
-
-
+        return change, per_layer_priority, np.array(global_deployment).tolist()

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,11 +233,28 @@ def unpack_update_batch(self, packed_update_info):
         ]
         return recovered
 
-    def get_expert_load(self) -> torch.Tensor:
-        load_info = self.shared_dict["load_info"]  # Tensor [L, W, local_experts_num]
-        logger.info(f"lt -- load_info {load_info=}...")
-        return load_info
-
+    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)
+        
+        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}...")

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

vllm_ascend/worker/model_runner_v1.py

@@ -1231,15 +1231,16 @@ def execute_model(
         intermediate_tensors: Optional[IntermediateTensors] = None,
     ) -> Union[ModelRunnerOutput, torch.Tensor]:
 
-        if self.dynamic_eplb:
-            self.eplb_updator.forward_before()
-
         with ProfileExecuteDuration().capture_async(
                 "prepare input and forward"):
             self._update_states(scheduler_output)
             if not scheduler_output.total_num_scheduled_tokens:
                 # Return empty ModelRunnerOuptut if there's no work to do.
                 return EMPTY_MODEL_RUNNER_OUTPUT
+
+            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,