[EPLB]Initializing eplb when dynamic eplb.

Author: offline0806

vllm_ascend/eplb/adaptor/vllm_adaptor.py

@@ -187,8 +187,6 @@ def _export_tensor_to_file(self, expert_maps, expert_map_record_path: str):
 
                 record["layer_list"].append(layer_record)
 
-            print(record)
-
             with open(expert_map_record_path, "w") as f:
                 json.dump(record, f, indent=4)
 
@@ -201,7 +199,6 @@ def do_update_expert_weight(self, layer_id, local_expert_to_replace,
         for expert_tensor, buffer_tensor in zip(
                 self.expert_param_per_layer[layer_id][local_expert_to_replace],
                 self.buffer_tensor_list[buffer_tensor_id]):
-            # expert_tensor.copy_(buffer_tensor)
             expert_tensor = buffer_tensor.clone()
 
     def do_update_log2phy_map(self, layer_id, updated_log2phy_map):

vllm_ascend/ops/common_fused_moe.py

@@ -27,17 +27,17 @@
     FusedMoEParallelConfig  # isort: skip
 from vllm.model_executor.layers.fused_moe.layer import (
     FusedMoE, UnquantizedFusedMoEMethod, determine_expert_map)
-
+from vllm_ascend.eplb.core.eplb_utils import (
+    determine_default_expert_map,
+    determine_default_log2phy_map)
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.distributed.parallel_state import get_mc2_group
 from vllm_ascend.ops.expert_load_balancer import ExpertLoadBalancer
 from vllm_ascend.ops.moe.experts_selector import select_experts
 from vllm_ascend.ops.moe.moe_comm_method import (AllGatherCommImpl,
                                                  AlltoAllCommImpl, MC2CommImpl)
 from vllm_ascend.ops.moe.token_dispatcher import setup_token_dispatchers
-from vllm_ascend.utils import ACL_FORMAT_FRACTAL_NZ, is_310p
 from vllm_ascend.utils import ACL_FORMAT_FRACTAL_NZ, is_310p, vllm_version_is
-from vllm.logger import logger
 
 original_unquantized_fused_moe_init_func = UnquantizedFusedMoEMethod.__init__
 
@@ -298,31 +298,26 @@ def __init__(
         self.moe_config.mc2_group = get_mc2_group()
         ascend_config = get_ascend_config()
         self.dynamic_eplb = ascend_config.dynamic_eplb
-        self.expert_map_path = ascend_config.expert_map_path
-        self.global_redundant_expert_num = ascend_config.init_redundancy_expert
-        self.global_num_experts = num_experts + self.global_redundant_expert_num
-        if self.expert_map_path and os.path.exists(self.expert_map_path) and os.access(self.expert_map_path, os.R_OK):
-            self.expert_load_balancer = ExpertLoadBalancer(self.expert_map_path, self.global_num_experts)
-            self.local_num_experts, self.expert_map = (self.expert_load_balancer.get_rank_placement_map(self.moe_instance_id, self.ep_rank))
-            self.log2phy = self.expert_load_balancer.get_rank_log2phy_map(self.moe_instance_id, self.ep_rank).npu()
-            self.global_redundant_expert_num = (self.expert_load_balancer.get_global_redundant_expert_num())
-        else:
-            self.local_num_experts, self.expert_map = determine_expert_map(self.ep_size, self.ep_rank, self.global_num_experts)
-            if self.dynamic_eplb:
-                self.global_redundant_expert_num = ascend_config.init_redundancy_expert
-                from vllm_ascend.eplb.core.eplb_utils import (
-                    determine_default_expert_map,
-                    determine_default_log2phy_map)
-                self.local_num_experts, self.expert_map = determine_default_expert_map(
-                    self.global_num_experts, self.ep_size, self.ep_rank,
-                    self.global_redundant_expert_num)
-                self.log2phy = determine_default_log2phy_map(
-                    self.global_num_experts, self.ep_size, self.ep_rank,
-                    self.global_redundant_expert_num)
-
-        self.moe_load = None
-        local_num_experts = (torch.sum(self.expert_map != -1) if self.expert_map is not None else num_experts)
         if self.dynamic_eplb:
+            self.expert_map_path = ascend_config.expert_map_path
+            self.global_redundant_expert_num = ascend_config.init_redundancy_expert
+            self.global_num_experts = num_experts + self.global_redundant_expert_num
+            if self.expert_map_path and os.path.exists(self.expert_map_path) and os.access(self.expert_map_path, os.R_OK):
+                self.expert_load_balancer = ExpertLoadBalancer(self.expert_map_path, self.global_num_experts)
+                self.local_num_experts, self.expert_map = (self.expert_load_balancer.get_rank_placement_map(self.moe_instance_id, self.ep_rank))
+                self.log2phy = self.expert_load_balancer.get_rank_log2phy_map(self.moe_instance_id, self.ep_rank).npu()
+                self.global_redundant_expert_num = (self.expert_load_balancer.get_global_redundant_expert_num())
+            else:
+                self.local_num_experts, self.expert_map = determine_expert_map(self.ep_size, self.ep_rank, self.global_num_experts)
+                if self.dynamic_eplb:
+                    self.global_redundant_expert_num = ascend_config.init_redundancy_expert
+                    self.local_num_experts, self.expert_map = determine_default_expert_map(
+                        self.global_num_experts, self.ep_size, self.ep_rank,
+                        self.global_redundant_expert_num)
+                    self.log2phy = determine_default_log2phy_map(
+                        self.global_num_experts, self.ep_size, self.ep_rank,
+                        self.global_redundant_expert_num)
+            local_num_experts = (torch.sum(self.expert_map != -1) if self.expert_map is not None else num_experts)
             self.moe_load = torch.zeros(local_num_experts, dtype=torch.int64)
 
 

vllm_ascend/ops/fused_moe.py

@@ -37,6 +37,9 @@
     FusedMoE, UnquantizedFusedMoEMethod, determine_expert_map)
 from vllm.model_executor.layers.quantization.base_config import \
     QuantizationConfig
+from vllm_ascend.eplb.core.eplb_utils import (
+    determine_default_expert_map,
+    determine_default_log2phy_map)
 
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ascend_forward_context import FusedMoEState
@@ -290,37 +293,29 @@ def __init__(
             self.moe_parallel_config.ep_size, is_deepseek_v3_r1)
 
         ascend_config = get_ascend_config()
-        expert_map_path = ascend_config.expert_map_path
         self.dynamic_eplb = ascend_config.dynamic_eplb
-        if expert_map_path and os.path.exists(expert_map_path):
-            # moe expert load balance
-            expert_load_balancer = ExpertLoadBalancer(expert_map_path,
-                                                      self.global_num_experts)
-            self.local_num_experts, self.expert_map = \
-                                expert_load_balancer.get_rank_placement_map(
-                                                self.moe_instance_id,
-                                                get_ep_group().rank_in_group)
-            self.log2phy = expert_load_balancer.get_rank_log2phy_map(
-                self.moe_instance_id,
-                get_ep_group().rank_in_group)
-            self.global_redundant_expert_num = \
-                        expert_load_balancer.get_global_redundant_expert_num()
-        else:
-            # Create a tensor of size num_experts filled with -1
-            self.local_num_experts, self.expert_map = determine_expert_map(
-                self.ep_size,
-                get_ep_group().rank_in_group, self.global_num_experts)
-            if self.dynamic_eplb:
-                self.global_redundant_expert_num = ascend_config.init_redundancy_expert
-                from vllm_ascend.eplb.core.eplb_utils import (
-                    determine_default_expert_map,
-                    determine_default_log2phy_map)
-                self.local_num_experts, self.expert_map = determine_default_expert_map(
-                    self.global_num_experts, self.ep_size, self.ep_rank,
-                    self.global_redundant_expert_num)
-                self.log2phy = determine_default_log2phy_map(
-                    self.global_num_experts, self.ep_size, self.ep_rank,
-                    self.global_redundant_expert_num)
+        if self.dynamic_eplb:
+            self.expert_map_path = ascend_config.expert_map_path
+            self.global_redundant_expert_num = ascend_config.init_redundancy_expert
+            self.global_num_experts = num_experts + self.global_redundant_expert_num
+            if self.expert_map_path and os.path.exists(self.expert_map_path) and os.access(self.expert_map_path, os.R_OK):
+                self.expert_load_balancer = ExpertLoadBalancer(self.expert_map_path, self.global_num_experts)
+                self.local_num_experts, self.expert_map = (self.expert_load_balancer.get_rank_placement_map(self.moe_instance_id, self.ep_rank))
+                self.log2phy = self.expert_load_balancer.get_rank_log2phy_map(self.moe_instance_id, self.ep_rank).npu()
+                self.global_redundant_expert_num = (self.expert_load_balancer.get_global_redundant_expert_num())
+            else:
+                self.local_num_experts, self.expert_map = determine_expert_map(self.ep_size, self.ep_rank, self.global_num_experts)
+                if self.dynamic_eplb:
+                    self.global_redundant_expert_num = ascend_config.init_redundancy_expert
+                    self.local_num_experts, self.expert_map = determine_default_expert_map(
+                        self.global_num_experts, self.ep_size, self.ep_rank,
+                        self.global_redundant_expert_num)
+                    self.log2phy = determine_default_log2phy_map(
+                        self.global_num_experts, self.ep_size, self.ep_rank,
+                        self.global_redundant_expert_num)
+            local_num_experts = (torch.sum(self.expert_map != -1) if self.expert_map is not None else num_experts)
+            self.moe_load = torch.zeros(local_num_experts, dtype=torch.int64)
+
 
         self.enable_shared_expert_dp = ascend_config.enable_shared_expert_dp
 

vllm_ascend/torchair/ops/torchair_fused_moe.py

@@ -37,7 +37,9 @@
     FusedMoE, UnquantizedFusedMoEMethod, determine_expert_map)
 from vllm.model_executor.layers.quantization.base_config import \
     QuantizationConfig
-
+from vllm_ascend.eplb.core.eplb_utils import (
+    determine_default_expert_map,
+    determine_default_log2phy_map)
 from vllm_ascend.ascend_config import get_ascend_config
 from vllm_ascend.ascend_forward_context import FusedMoEState
 from vllm_ascend.distributed.parallel_state import get_mc2_group
@@ -1012,27 +1014,27 @@ def __init__(
 
         ascend_config = get_ascend_config()
         self.dynamic_eplb = ascend_config.dynamic_eplb
-        self.expert_map_path = ascend_config.expert_map_path
-        self.global_redundant_expert_num = ascend_config.init_redundancy_expert
-        self.global_num_experts = num_experts + self.global_redundant_expert_num
-        if self.expert_map_path and os.path.exists(self.expert_map_path):
-            self.expert_load_balancer = ExpertLoadBalancer(self.expert_map_path, self.global_num_experts)
-            self.local_num_experts, self.expert_map = (self.expert_load_balancer.get_rank_placement_map(self.moe_instance_id, self.ep_rank))
-            self.log2phy = self.expert_load_balancer.get_rank_log2phy_map(self.moe_instance_id, self.ep_rank).npu()
-            self.global_redundant_expert_num = (self.expert_load_balancer.get_global_redundant_expert_num())
-        else:
-            self.local_num_experts, self.expert_map = determine_expert_map(self.ep_size, self.ep_rank, self.global_num_experts)
-            if self.dynamic_eplb:
-                self.global_redundant_expert_num = ascend_config.init_redundancy_expert
-                from vllm_ascend.eplb.core.eplb_utils import (
-                    determine_default_expert_map,
-                    determine_default_log2phy_map)
-                self.local_num_experts, self.expert_map = determine_default_expert_map(
-                    self.global_num_experts, self.ep_size, self.ep_rank,
-                    self.global_redundant_expert_num)
-                self.log2phy = determine_default_log2phy_map(
-                    self.global_num_experts, self.ep_size, self.ep_rank,
-                    self.global_redundant_expert_num)
+        if self.dynamic_eplb:
+            self.expert_map_path = ascend_config.expert_map_path
+            self.global_redundant_expert_num = ascend_config.init_redundancy_expert
+            self.global_num_experts = num_experts + self.global_redundant_expert_num
+            if self.expert_map_path and os.path.exists(self.expert_map_path) and os.access(self.expert_map_path, os.R_OK):
+                self.expert_load_balancer = ExpertLoadBalancer(self.expert_map_path, self.global_num_experts)
+                self.local_num_experts, self.expert_map = (self.expert_load_balancer.get_rank_placement_map(self.moe_instance_id, self.ep_rank))
+                self.log2phy = self.expert_load_balancer.get_rank_log2phy_map(self.moe_instance_id, self.ep_rank).npu()
+                self.global_redundant_expert_num = (self.expert_load_balancer.get_global_redundant_expert_num())
+            else:
+                self.local_num_experts, self.expert_map = determine_expert_map(self.ep_size, self.ep_rank, self.global_num_experts)
+                if self.dynamic_eplb:
+                    self.global_redundant_expert_num = ascend_config.init_redundancy_expert
+                    self.local_num_experts, self.expert_map = determine_default_expert_map(
+                        self.global_num_experts, self.ep_size, self.ep_rank,
+                        self.global_redundant_expert_num)
+                    self.log2phy = determine_default_log2phy_map(
+                        self.global_num_experts, self.ep_size, self.ep_rank,
+                        self.global_redundant_expert_num)
+            local_num_experts = (torch.sum(self.expert_map != -1) if self.expert_map is not None else num_experts)
+            self.moe_load = torch.zeros(local_num_experts, dtype=torch.int64)
 
         self.torchair_graph_enabled = ascend_config.torchair_graph_config.enabled
         self.enable_multistream_moe = \