merge from remote main

Author: lt

vllm_ascend/eplb/adaptor/vllm_adaptor.py

@@ -38,7 +38,7 @@ def __init__(self, model, **args):
         self.num_moe_layers = self.model.config.num_hidden_layers - self.num_dense_layers
         self.global_expert_num = self.model.config.n_routed_experts
 
-        
+
         # TODO: init self.expert_weight_names depending on different model types, only deepseek v3 w8a8 is supported here
         self.expert_weight_names = ["w13_weight", "w2_weight", "w13_weight_scale", "w13_weight_offset",
             "w2_weight_scale", "w2_weight_offset"]
@@ -62,6 +62,8 @@ def __init__(self, model, **args):
             self.log2phy_map_per_layer[self.num_dense_layers + layer_idx] =\
                 self.model.get_log2phy_map(self.num_dense_layers + layer_idx)
 
+        self.all_topk_ids = []
+
     def init_buffer_tensor(self, num_buffer_tensor):
         for name in self.expert_weight_names:
             complete_name = "model.layers." + str(self.num_dense_layers) + ".mlp.experts." + name
@@ -82,39 +84,43 @@ def init_expert_param_per_layer(self):
                         for name in self.expert_weight_names]
                 )
 
-    def get_rank_expert_workload(
-        self,
-        num_moe_layers: int,
-        dummy_run = False
-    ) -> torch.Tensor:
-
-        all_topk_ids = [self.model.get_topk_ids(i) for i in range(num_moe_layers)]
-        stacked = torch.stack(all_topk_ids, dim=0)      
-        L, B, K = stacked.shape
-        N = B * K
-        device = stacked.device
-        G = self.global_expert_num
-
-        if not hasattr(self, "cum_moe_load") or self.cum_moe_load is None:
-            self.cum_moe_load = torch.zeros((L, G),
-                                            dtype=torch.int64,
-                                            device=device)
-
+    def collect_topk_ids(self, dummy_run=False):
         if dummy_run:
-            return self.cum_moe_load
-
-        ids1d = stacked.view(-1).to(torch.int64)       
-
-        row_idx = torch.arange(L, device=device).repeat_interleave(N) 
-
-        combined = row_idx * G + ids1d                     
-
-        counts = torch.bincount(combined, minlength=L * G)  
-        workload = counts.view(L, G)                       
-
-        self.cum_moe_load.add_(workload)
-    
-        return self.cum_moe_load
+            return 
+        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:     
+            for l in range(self.num_moe_layers):  
+                t = period_data[l]       
+                flat_list_per_layer[l].append(t.reshape(-1))  
+
+        index_2d = torch.nn.utils.rnn.pad_sequence(
+            [torch.cat(flat_list_per_layer[l]) for l in range(self.num_moe_layers)],
+            batch_first=True, padding_value=-1     
+        ).to(device)               
+
+        mask = index_2d != -1
+        index_2d = index_2d.masked_select(mask).reshape(self.num_moe_layers, -1)
+        src_2d   = torch.ones_like(index_2d, dtype=torch.int64)
+
+        self.moe_load.scatter_add_(dim=1, index=index_2d, src=src_2d)
+
+        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)
@@ -136,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
@@ -57,7 +61,7 @@ def init_eplb(self, expert_map_path):
         self.cur_iterations: torch.int64 = 0
 
         self.wait_worker_iterations: torch.int64 = 0
-        self.num_wait_worker_iterations: torch.int64 = 10
+        self.num_wait_worker_iterations: torch.int64 = 20
 
         self.planner_block_queue = Queue()
         self.block_update_queue = Queue(maxsize=1)
@@ -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,30 +131,31 @@ 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.get_rank_expert_workload(self.num_moe_layers,dummy_run)
-            if not self.update_in_flight:
-                load_gather_iteration, update_iteration = self.get_update_iteration()
-                if load_gather_iteration:
-                    moe_load = self.compute_and_set_moe_load(dummy_run)
-                if update_iteration:
-                    self.wakeup_eplb_worker()
-                    self.update_in_flight = True
-                    self.wait_worker_iterations = 0
-                    self.weight_loading = False
-
-            if self.update_in_flight:
-                self.wait_worker_iterations = self.wait_worker_iterations + 1
-
-            self.eplb_loader.update_expert_map_and_weight(self.reqs, self.redundant_enable)
+        self.adaptor.collect_topk_ids(dummy_run)
+        if not self.update_in_flight:
+            load_gather_iteration, update_iteration = self.get_update_iteration()
+            if load_gather_iteration:
+                moe_load = self.compute_and_set_moe_load()
+            if update_iteration:
+                self.wakeup_eplb_worker()
+                self.update_in_flight = True
+                self.wait_worker_iterations = 0
+                self.weight_loading = False
+
+        if self.update_in_flight:
+            self.wait_worker_iterations = self.wait_worker_iterations + 1
+
+        self.eplb_loader.update_expert_map_and_weight(self.reqs, self.redundant_enable)
 
     def compute_and_set_moe_load(self,dummy_run=False):
-        local_load = self.adaptor.get_rank_expert_workload(self.num_moe_layers,dummy_run)
+        local_load = self.adaptor.get_rank_expert_workload()
+
         self._gather_buffer = None
         if dist.is_initialized():
             self.world_size = dist.get_world_size()
@@ -173,7 +180,7 @@ def compute_and_set_moe_load(self,dummy_run=False):
     def warm_up_eplb(self):
 
         self.get_init_expert_map()
-        
+        self.adaptor.collect_topk_ids(dummy_run=False)
         self.compute_and_set_moe_load()
 
         src_tensor = torch.empty((1,), device=self.device)
@@ -228,29 +235,18 @@ 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
+    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 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.

vllm_ascend/models/deepseek_v2.py

@@ -772,5 +772,12 @@ def get_all_expert_map(self,num_moe_layers):
     def get_topk_ids(self,layer_id):
         return self.model.layers[layer_id+3].mlp.experts.topk_ids
 
+    def get_all_topk_ids(self,num_moe_layers):
+        all_topk_id = []
+        for layer_id in range(num_moe_layers):
+            load_tensor = self.get_topk_ids(layer_id)  
+            all_topk_id.append(load_tensor)
+        return all_topk_id
+
 class CustomDeepseekV3ForCausalLM(CustomDeepseekV2ForCausalLM):
     pass

vllm_ascend/quantization/w8a8_dynamic.py

@@ -678,7 +678,7 @@ def apply(
                 global_redundant_expert_num=global_redundant_expert_num,
                 shared_experts=shared_experts,
                 **kwargs), topk_ids
-        elif self.torchair_graph_enabled or self.ep_group.world_size == 1:
+        elif fused_moe_state == FusedMoEState.AllGather:
             return fused_experts(hidden_states=x,
                                  w1=layer.w13_weight,
                                  w1_scale=layer.w13_weight_scale,