Merge pull request #107 from raindaywhu/dev_whq_eplb2

modify serialization of eplb process

Author: wanghanqingLYT

vllm_ascend/eplb/core/worker/eplb_worker.py

@@ -336,30 +336,22 @@ def pack_update_info(self, update_info_generator):
 
         for send_info, recv_info, new_expert_map, layer_id in update_info_generator:
 
-            send_all.append(send_info)
-            recv_all.append(recv_info)
+            send_info_this_rank = send_info[self.rank_id] if self.rank_id in send_info else []
+            recv_info_this_rank = recv_info[self.rank_id] if self.rank_id in recv_info else []
+            send_all.append(send_info_this_rank)
+            recv_all.append(recv_info_this_rank)
 
-            maps.append(new_expert_map[self.rank_id])
+            maps.append(new_expert_map[self.rank_id].numpy().tolist())
 
-            if self.redundant_enable is not None:
+            if self.redundant_enable:
                 log2phy_map = ExpertMapUtils.generate_log2phy_map(new_expert_map)
-                log2phy_all.append(log2phy_map)
+                log2phy_all.append(log2phy_map[self.rank_id].numpy().tolist())
+            else:
+                log2phy_all.append([])
 
             layer_ids.append(layer_id)
 
-        # 把 list of Tensor 堆成一个大 Tensor
-        stacked_maps      = torch.stack(maps,      dim=0)
-        layer_id_tensor   = torch.as_tensor(layer_ids, dtype=torch.int64)
-        stacked_maps.share_memory_()
-        layer_id_tensor.share_memory_()
-
-        if self.redundant_enable:
-            stacked_log2phy = torch.stack(log2phy_all, dim=0)
-            stacked_log2phy.share_memory_()
-        else:
-            stacked_log2phy = None
-
-        return send_all, recv_all, stacked_maps, stacked_log2phy, layer_id_tensor
+        return list(zip(send_all, recv_all, maps, log2phy_all, layer_ids))
 
 class EplbProcess:
     def __init__(self, shared_dict, planner_q, block_update_q, redundant_enable, policy_type: int = 0, enable_d2d: bool = True):

vllm_ascend/eplb/eplb_updator.py

@@ -16,6 +16,7 @@
 #
 
 import torch
+import numpy
 from typing import Dict, List
 import torch.distributed as dist
 import vllm.envs as envs
@@ -111,19 +112,19 @@ def forward_before(self):
         # Batch after eplb process being triggered, get update info provided by eplb process
         if self.update_in_flight and self.weight_update_counter == 0 and self.wait_worker_iterations == self.num_wait_worker_iterations:
             self.wait_worker_iterations = 0
-            packed_update_info = self.block_update_queue.get()
-            self.update_info_all = self.unpack_update_batch(packed_update_info)
+            self.update_info_all = self.block_update_queue.get()
             self.weight_loading = True
 
         if self.update_in_flight and self.weight_loading and self.weight_update_counter < self.num_moe_layers:
             (expert_send_info, expert_recv_info, updated_expert_map, log2phy_map, layer_id) = self.update_info_all.pop(0)
             rank_id = torch.distributed.get_rank()
-            self.eplb_loader.set_log2phy_map(log2phy_map)
-            expert_send_info_this_rank = expert_send_info[rank_id] if rank_id in expert_send_info else []
-            expert_recv_info_this_rank = expert_recv_info[rank_id] if rank_id in expert_recv_info else []
+            if self.redundant_enable:
+                log2phy_map_this_rank = torch.from_numpy(numpy.array(log2phy_map))
+                self.eplb_loader.set_log2phy_map(log2phy_map_this_rank)
+            updated_expert_map_this_rank = torch.from_numpy(numpy.array(updated_expert_map))
             #logger.info(f"check update info, layer = {layer_id}, send = {expert_send_info_this_rank}, recv = {expert_recv_info_this_rank}")
-            self.eplb_loader.generate_expert_d2d_transfer_task(expert_send_info_this_rank,
-                expert_recv_info_this_rank, updated_expert_map, layer_id + 3)
+            self.eplb_loader.generate_expert_d2d_transfer_task(expert_send_info, expert_recv_info,
+                updated_expert_map_this_rank, layer_id + self.adaptor.num_dense_layers)
             self.weight_update_counter += 1
             if self.weight_update_counter == self.num_moe_layers:
                 self.weight_update_counter = 0