fix nz

Author: raindaywhu

vllm_ascend/eplb/adaptor/vllm_adaptor.py

@@ -69,15 +69,16 @@ def __init__(self, model, **args):
         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
-            expert_tensor = self.param_dict[complete_name].data[
-                0:num_buffer_tensor]
-            buffer_tensors = torch.empty_like(expert_tensor)
-            for buffer_id in range(num_buffer_tensor):
+        for buffer_id in range(num_buffer_tensor):
+            for name in self.expert_weight_names:
+                complete_name = "model.layers." + str(
+                    self.num_dense_layers) + ".mlp.experts." + name
+                expert_tensor = self.param_dict[complete_name].data[0]
+                if name in ["w13_weight", "w2_weight"]:
+                    expert_tensor = expert_tensor.clone()
+                buffer_tensor = torch.empty_like(expert_tensor)
                 self.buffer_tensor_list[buffer_id].append(
-                    buffer_tensors[buffer_id])
+                    buffer_tensor)
 
     def init_expert_param_per_layer(self):
         num_local_expert = self.param_dict["model.layers." + str(self.num_dense_layers) + \

vllm_ascend/eplb/core/eplb_device_transfer_loader.py

@@ -18,6 +18,7 @@
 
 import torch.distributed as dist
 from vllm.logger import logger
+from vllm_ascend.ascend_config import get_ascend_config
 
 
 class ExpertWeightUpdateState(Enum):
@@ -38,6 +39,7 @@ def __init__(self, eplb_adaptor):
         self.state = ExpertWeightUpdateState.WAITING
         self.recv_expert_list = []
         self.mock_flag = True
+        self.enable_weight_nz_layout = get_ascend_config().enable_weight_nz_layout
 
     def generate_expert_d2d_transfer_task(self, expert_send_info,
                                           expert_recv_info, updated_expert_map,
@@ -61,10 +63,14 @@ def generate_expert_d2d_transfer_task(self, expert_send_info,
             dst_rank, global_expert_id_to_send = send_info
             local_expert_id = self.eplb_adaptor.expert_map_per_layer_cpu[
                 layer_id][global_expert_id_to_send].item()
+            idx = 0
             for src_tensor in self.eplb_adaptor.expert_param_per_layer[
-                    layer_id][local_expert_id]:
+                layer_id][local_expert_id]:
+                if self.enable_weight_nz_layout and idx < 2:
+                    src_tensor = src_tensor.clone()
                 self.comm_op_list.append(
                     dist.P2POp(dist.isend, src_tensor, dst_rank))
+                idx += 1
 
         buffer_tensor_id = 0
         for recv_info in expert_recv_info: