EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

337
vllm_ascend/eplb/adaptor/vllm_adaptor.py
View File

@@ -23,169 +23,154 @@ import torch.distributed as dist
from vllm.logger import logger
from vllm_ascend.ascend_config import get_ascend_config
from vllm_ascend.eplb.adaptor.abstract_adaptor import EplbAdaptor
from vllm_ascend.quantization.quant_type import QuantType
EPLB_EXPERT_WEIGHT_NAMES = {
(QuantType.NONE, False): ("w13_weight", "w2_weight"),
(QuantType.NONE, True): ("w13_weight_list", "w2_weight_list"),
(QuantType.W8A8, False): (
"w13_weight_list",
"w2_weight_list",
"w13_weight_scale_fp32_list",
"w2_weight_scale_list",
),
(QuantType.W8A8, True): (
"w13_weight_list",
"w2_weight_list",
"w13_weight_scale_fp32_list",
"w2_weight_scale_list",
"fused_w1_scale_list",
"fused_w2_scale_list",
),
(QuantType.W4A8, True): (
"w13_weight_list",
"w2_weight_list",
"w13_weight_scale_list",
"w2_weight_scale_list",
"w13_scale_bias_list",
"w2_scale_bias_list",
),
(QuantType.MXFP4, False): ("w13_weight", "w2_weight", "w13_weight_scale", "w2_weight_scale"),
(QuantType.MXFP4, True): ("w13_weight", "w2_weight", "w13_weight_scale", "w2_weight_scale"),
(QuantType.MXFP8, False): ("w13_weight", "w2_weight", "w13_weight_scale", "w2_weight_scale"),
(QuantType.MXFP8, True): ("w13_weight", "w2_weight", "w13_weight_scale", "w2_weight_scale"),
}
class VllmEplbAdaptor(EplbAdaptor):
class VllmEplbAdaptor:
_registered_moe_layers: list["torch.nn.Module"] = []
@staticmethod
def register_layer(layer: "torch.nn.Module") -> None:
"""Register a MoE layer for EPLB. Called during layer initialization.
Only real layers call this; PPMissingLayer won't, so the registry
naturally contains only layers on this PP rank.
"""
VllmEplbAdaptor._registered_moe_layers.append(layer)
def __init__(self, model, **args):
super().__init__(**args)
self.model = model
if hasattr(model, "language_model"):
self.model = model.language_model
self.config = model.config.text_config
else:
self.model = model
self.config = model.config
self.rank_id = dist.get_rank()
self.world_size = dist.get_world_size()
self.param_dict = dict(self.model.named_parameters())
if self.model.config.model_type == "qwen3_moe":
self.num_dense_layers = 0
self.global_expert_num = self.model.config.num_experts
else:
self.num_dense_layers = self.model.config.first_k_dense_replace
self.global_expert_num = self.model.config.n_routed_experts
self.num_moe_layers = self.model.config.num_hidden_layers - self.num_dense_layers
self.init_redundancy_expert = get_ascend_config(
).init_redundancy_expert
self.num_dense_layers = getattr(self.config, "first_k_dense_replace", 0)
# TODO: init self.expert_weight_names depending on different model types, only deepseek v3 w8a8 and qwen3-moe is supported here
if self.model.quant_config is not None:
self.expert_weight_names = [
"w13_weight", "w2_weight", "w13_weight_scale",
"w13_weight_offset", "w2_weight_scale", "w2_weight_offset"
]
else:
self.expert_weight_names = ["w13_weight", "w2_weight"]
self.moe_layers = VllmEplbAdaptor._registered_moe_layers
self.num_moe_layers = len(self.moe_layers)
self.expert_map_per_layer = dict(
) # reference to expert map on device for expert map update
self.expert_map_per_layer_cpu = dict(
) # copy of expert map on CPU to avoid device synchronize frequently
for layer_idx in range(self.num_moe_layers):
self.expert_map_per_layer[self.num_dense_layers + layer_idx] = \
self.model.get_expert_map(self.num_dense_layers + layer_idx)
self.expert_map_per_layer_cpu = dict() # copy of expert map on CPU to avoid device synchronize frequently
# TODO: here we set number of buffer tensor equal to number of expert in each laryer, which can be improved
num_buffer_tensor = torch.where(
self.expert_map_per_layer[self.num_dense_layers] != -1)[0].numel()
self.buffer_tensor_list: list[list[Any]] = [
[] for _ in range(num_buffer_tensor)
]
self.init_buffer_tensor(num_buffer_tensor)
# Get num_local_experts from first real MoE layer
first_layer = self.moe_layers[0]
self.num_local_experts = first_layer.local_num_experts
self.ep_rank = first_layer.ep_rank
self.expert_param_per_layer = dict()
self.expert_weight_key_per_layer = dict()
self.init_expert_param_per_layer()
self.log2phy_map_per_layer = dict()
for layer_idx in range(self.num_moe_layers):
self.log2phy_map_per_layer[self.num_dense_layers + layer_idx] = \
self.model.get_log2phy_map(self.num_dense_layers + layer_idx)
num_buffer_tensor = self.num_local_experts
self.buffer_tensor_list: dict[Any, list[list[Any]]] = dict()
self.init_buffer_tensor(num_buffer_tensor)
self.all_topk_ids = []
self.log2phy_map_per_layer = dict()
for local_idx, layer in enumerate(self.moe_layers):
self.log2phy_map_per_layer[local_idx] = layer.get_log2phy_map()
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)
buffer_tensor_shapes: dict[Any, list[torch.Size]] = dict()
for local_idx, _ in enumerate(self.moe_layers):
expert_weight_key = self.expert_weight_key_per_layer[local_idx]
expert_weight_names = EPLB_EXPERT_WEIGHT_NAMES[expert_weight_key]
expert_tensors = [self.param_dict[f"{local_idx}.{name}"][0] for name in expert_weight_names]
expert_tensor_shapes = [tensor.shape for tensor in expert_tensors]
if expert_weight_key in self.buffer_tensor_list:
assert expert_tensor_shapes == buffer_tensor_shapes[expert_weight_key], (
f"EPLB expert weight shapes mismatch for {expert_weight_key}: "
f"expected {buffer_tensor_shapes[expert_weight_key]}, got {expert_tensor_shapes}"
)
continue
buffer_tensor_shapes[expert_weight_key] = expert_tensor_shapes
self.buffer_tensor_list[expert_weight_key] = [[] for _ in range(num_buffer_tensor)]
for buffer_id in range(num_buffer_tensor):
self.buffer_tensor_list[buffer_id].append(
buffer_tensors[buffer_id])
for expert_tensor in expert_tensors:
buffer_tensor = torch.empty_like(expert_tensor)
self.buffer_tensor_list[expert_weight_key][buffer_id].append(buffer_tensor)
def init_expert_param_per_layer(self):
num_local_expert = self.param_dict["model.layers." + str(self.num_dense_layers) + \
".mlp.experts." + self.expert_weight_names[0]].data.shape[0]
for moe_layer_id in range(self.num_moe_layers):
layer_idx = self.num_dense_layers + moe_layer_id
self.expert_param_per_layer[layer_idx] = list()
for local_expert_id in range(num_local_expert):
self.expert_param_per_layer[layer_idx].append([
self.param_dict["model.layers." + str(layer_idx) +
".mlp.experts." +
name].data[local_expert_id]
for name in self.expert_weight_names
])
self.param_dict = dict()
for local_idx, layer in enumerate(self.moe_layers):
quant_type = QuantType.NONE if self.model.quant_config is None else layer.quant_type
expert_weight_key = (quant_type, get_ascend_config().enable_fused_mc2 == 1)
if expert_weight_key[0] == QuantType.W4A8MXFP:
raise RuntimeError(f"EPLB not support {quant_type}")
if expert_weight_key not in EPLB_EXPERT_WEIGHT_NAMES:
raise ValueError(f"EPLB not support {quant_type} with fused MC2 {expert_weight_key[1]}")
expert_weight_names = EPLB_EXPERT_WEIGHT_NAMES[expert_weight_key]
self.expert_weight_key_per_layer[local_idx] = expert_weight_key
self.expert_param_per_layer[local_idx] = list()
for name in expert_weight_names:
param_key = f"{local_idx}.{name}"
self.param_dict[param_key] = getattr(layer, name)
for local_expert_id in range(self.num_local_experts):
per_expert_param = list()
for name in expert_weight_names:
per_expert_param.append(self.param_dict[f"{local_idx}.{name}"][local_expert_id])
self.expert_param_per_layer[local_idx].append(per_expert_param)
def get_rank_expert_workload(self) -> torch.Tensor:
self.moe_load = self.model.get_all_moe_loads()
loads = [layer.moe_load for layer in self.moe_layers]
self.moe_load = torch.stack(loads, dim=0) if loads else torch.empty(0)
return self.moe_load
def get_init_expert_map(self, num_moe_layers):
expert_map = self.model.get_all_expert_map(num_moe_layers)
if dist.is_initialized():
world_size = dist.get_world_size()
gathered = torch.empty(
(world_size, *expert_map.shape), # [W, L, E]
dtype=expert_map.dtype,
device=expert_map.device)
dist.all_gather_into_tensor(gathered, expert_map)
all_maps = gathered.permute(1, 0, 2)
all_expert_maps = all_maps.cpu()
for layer_idx in range(num_moe_layers):
self.expert_map_per_layer_cpu[self.num_dense_layers + layer_idx] = \
all_expert_maps[layer_idx][self.rank_id]
return all_expert_maps
def get_init_expert_map_from_file(self, num_moe_layers, expert_map_path):
try:
expert_map_tensor, layers_num, ranks_num = self._expert_file_to_tensor(
expert_map_path)
expert_map_all = self.local2global(expert_map_tensor)
except (TypeError, FileNotFoundError, OSError):
expert_map_all = self.determine_expert_map_all()
for layer_idx in range(num_moe_layers):
if self.model.config.model_type == "qwen3_moe":
self.expert_map_per_layer_cpu[layer_idx] = \
expert_map_all[layer_idx][self.rank_id]
else:
self.expert_map_per_layer_cpu[layer_idx + self.num_dense_layers] = \
expert_map_all[layer_idx][self.rank_id]
return expert_map_all
def _expert_file_to_tensor(self, expert_map_path: str):
with open(expert_map_path, "r") as f:
data = json.load(f)
layers_num = data["moe_layer_count"]
gpus_num = data["layer_list"][0]["device_count"]
tensor_data = []
for layer in data["layer_list"]:
device_data = []
for device in layer["device_list"]:
device_data.append(device["device_expert"])
tensor_data.append(device_data)
expert_map_tensor = torch.tensor(tensor_data, dtype=torch.int32)
return expert_map_tensor, layers_num, gpus_num
logger.error(f"failed to read expert_map_path: {expert_map_path}")
def clear_all_moe_loads(self):
for layer in self.moe_layers:
layer.clear_moe_load()
def _export_tensor_to_file(self, expert_maps, expert_map_record_path: str):
if self.rank_id == 0:
num_local_experts = expert_maps.max() + 1
expert_maps_local = self.global2local(expert_maps,
num_local_experts)
expert_maps_list = expert_maps_local.tolist()
record: dict[str, Any] = {
"moe_layer_count": len(expert_maps_list),
"layer_list": []
}
expert_maps_list = expert_maps.tolist()
record: dict[str, Any] = {"moe_layer_count": len(expert_maps_list), "layer_list": []}
for layer_idx, layer_data in enumerate(expert_maps_list):
layer_record: dict[str, Any] = {
"layer_id": layer_idx,
"device_count": len(layer_data),
"device_list": []
"device_list": [],
}
for device_idx, experts in enumerate(layer_data):
device_record = {
"device_id": device_idx,
"device_expert": experts
}
placement = [experts.index(i) for i in range(num_local_experts)]
device_record = {"device_id": device_idx, "device_expert": placement}
layer_record["device_list"].append(device_record)
record["layer_list"].append(layer_record)
@@ -194,96 +179,26 @@ class VllmEplbAdaptor(EplbAdaptor):
json.dump(record, f, indent=4)
def do_update_expert_map(self, layer_id, updated_expert_map):
self.expert_map_per_layer[layer_id] = updated_expert_map.clone()
self.expert_map_per_layer_cpu[layer_id] = updated_expert_map.clone()
self.expert_map_per_layer_cpu[layer_id].copy_(updated_expert_map)
def do_update_expert_weight(self, layer_id, local_expert_to_replace,
buffer_tensor_id):
def do_update_expert_weight(self, layer_id, local_expert_to_replace, buffer_tensor_id):
expert_weight_key = self.expert_weight_key_per_layer[layer_id]
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 = buffer_tensor.clone()
logger.debug(f"Expert tensor shape is :{expert_tensor.shape}")
self.expert_param_per_layer[layer_id][local_expert_to_replace],
self.buffer_tensor_list[expert_weight_key][buffer_tensor_id],
):
expert_tensor.copy_(buffer_tensor)
logger.debug("Expert tensor shape is :%s", expert_tensor.shape)
def do_update_log2phy_map(self, layer_id, updated_log2phy_map):
if self.log2phy_map_per_layer[layer_id] is not None:
self.log2phy_map_per_layer[layer_id].copy_(updated_log2phy_map)
def global2local(self, placement: torch.Tensor,
E_local: int) -> torch.Tensor:
def get_global_expert_map(self):
all_layer_global_expert_map = []
for local_idx, layer in enumerate(self.moe_layers):
map_cpu = layer.global_expert_map.cpu()
all_layer_global_expert_map.append(map_cpu)
self.expert_map_per_layer_cpu[local_idx] = map_cpu[self.ep_rank]
L, G, _ = placement.shape
device = placement.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]
pt_local[l_idx, g_idx, slot_idx] = k_idx
return pt_local
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 determine_expert_map_all(self):
if self.world_size == 1:
local_ids = torch.arange(self.global_expert_num, dtype=torch.int32)
return local_ids.view(1, 1, -1).expand(self.num_moe_layers, 1, -1)
local_num_experts = self.global_expert_num // self.world_size
expert_map_all = torch.full(
(self.num_moe_layers, self.world_size, self.global_expert_num),
-1,
dtype=torch.int32)
for r in range(self.world_size):
if r < self.world_size - 1:
start = r * local_num_experts
end = (r + 1) * local_num_experts
local_count = local_num_experts
else:
start = r * local_num_experts
end = self.global_expert_num
local_count = self.global_expert_num - r * local_num_experts
if r < self.init_redundancy_expert:
local_count += 1
if end < self.global_expert_num:
end += 1
else:
start -= 1
local_ids = torch.arange(local_count, dtype=torch.int32)
expert_map_all[:, r, start:end] = local_ids.unsqueeze(0).expand(
self.num_moe_layers, -1)
return expert_map_all
return torch.stack(all_layer_global_expert_map)

85
vllm_ascend/eplb/core/eplb_device_transfer_loader.py
View File

@@ -18,6 +18,9 @@ from enum import Enum
import torch.distributed as dist
from vllm.logger import logger
from vllm.v1.utils import record_function_or_nullcontext
from vllm_ascend.distributed.parallel_state import get_dynamic_eplb_group
class ExpertWeightUpdateState(Enum):
@@ -27,7 +30,6 @@ class ExpertWeightUpdateState(Enum):
class D2DExpertWeightLoader:
def __init__(self):
self.comm_op_list = None
self.updated_expert_map = None
@@ -35,50 +37,45 @@ class D2DExpertWeightLoader:
self.layer_id = -1 # layer id to be updated
self.state = ExpertWeightUpdateState.WAITING
self.recv_expert_list = []
self.mock_flag = True
self.num_layers = 0
self.comm_group = get_dynamic_eplb_group()
def set_adator(self, eplb_adaptor):
self.eplb_adaptor = eplb_adaptor
def generate_expert_d2d_transfer_task(self, expert_send_info,
expert_recv_info, updated_expert_map,
layer_id):
def generate_expert_d2d_transfer_task(self, expert_send_info, expert_recv_info, updated_expert_map, layer_id):
# When current send/recv and weight.expert_map update tasks are not finished, cannot accept new d2d task
if self.state != ExpertWeightUpdateState.WAITING:
logger.error(
"current d2d weight update tasks are on-going, cannot accept new weight update task"
logger.warning_once(
"[eplb/d2d_loader] Current D2D weight update is on-going, cannot accept new update task"
)
return
# If neither send nor receive task is needed for this layer on this rank, return
if not (expert_send_info or expert_recv_info):
return
self.updated_expert_map = updated_expert_map
self.layer_id = layer_id
self.comm_op_list = []
for send_info in 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()
for src_tensor in self.eplb_adaptor.expert_param_per_layer[
layer_id][local_expert_id]:
local_expert_id = self.eplb_adaptor.expert_map_per_layer_cpu[layer_id][global_expert_id_to_send].item()
for src_tensor in self.eplb_adaptor.expert_param_per_layer[layer_id][local_expert_id]:
self.comm_op_list.append(
dist.P2POp(dist.isend, src_tensor, dst_rank))
dist.P2POp(
dist.isend, src_tensor, self.comm_group.ranks[dst_rank], group=self.comm_group.device_group
)
)
buffer_tensor_id = 0
for recv_info in expert_recv_info:
for buffer_tensor_id, recv_info in enumerate(expert_recv_info):
recv_rank, global_expert_id_to_recv = recv_info
for buffer_tensor in self.eplb_adaptor.buffer_tensor_list[
buffer_tensor_id]:
expert_weight_key = self.eplb_adaptor.expert_weight_key_per_layer[layer_id]
for buffer_tensor in self.eplb_adaptor.buffer_tensor_list[expert_weight_key][buffer_tensor_id]:
self.comm_op_list.append(
dist.P2POp(dist.irecv, buffer_tensor, recv_rank))
local_expert_to_replace = self.updated_expert_map[
global_expert_id_to_recv].item()
self.recv_expert_list.append(
(local_expert_to_replace, buffer_tensor_id))
buffer_tensor_id += 1
dist.P2POp(
dist.irecv, buffer_tensor, self.comm_group.ranks[recv_rank], group=self.comm_group.device_group
)
)
local_expert_to_replace = self.updated_expert_map[global_expert_id_to_recv].item()
self.recv_expert_list.append((local_expert_to_replace, buffer_tensor_id))
self.state = ExpertWeightUpdateState.READY
@@ -86,7 +83,7 @@ class D2DExpertWeightLoader:
self.updated_log2phy_map = log2phy_map
def asyn_expert_weight_transfer(self, reqs):
# Only when send/recv tasks are parsed into self.comm_op_list, d2d send/recv tasks can be luanched
# Only when send/recv tasks are parsed into self.comm_op_list, d2d send/recv tasks can be launched
if self.state != ExpertWeightUpdateState.READY:
return
@@ -98,40 +95,44 @@ class D2DExpertWeightLoader:
self.state = ExpertWeightUpdateState.TRANSFERRING
def update_expert_map_and_weight(self, reqs):
# Only after send/recv tasks have been luanched, expert_map and weight can be updated
# Only after send/recv tasks have been launched, expert_map and weight can be updated
if self.state != ExpertWeightUpdateState.TRANSFERRING:
return
# Waiting for send/recv tasks finish
for req in reqs:
req.wait()
if reqs:
with record_function_or_nullcontext("EPLB weight D2D wait"):
for req in reqs:
req.wait()
if self.comm_op_list is not None:
self.comm_op_list = None
# update expert_map
self.eplb_adaptor.do_update_expert_map(self.layer_id,
self.updated_expert_map)
self.eplb_adaptor.do_update_expert_map(self.layer_id, self.updated_expert_map)
# update log2phy_map
self.eplb_adaptor.do_update_log2phy_map(self.layer_id,
self.updated_log2phy_map)
self.eplb_adaptor.do_update_log2phy_map(self.layer_id, self.updated_log2phy_map)
# update expert weight
buffer_tensor_id = 0
for recv_expert_info in self.recv_expert_list:
local_expert_to_replace, buffer_tensor_id = recv_expert_info
self.eplb_adaptor.do_update_expert_weight(self.layer_id,
local_expert_to_replace,
buffer_tensor_id)
self.eplb_adaptor.do_update_expert_weight(self.layer_id, local_expert_to_replace, buffer_tensor_id)
logger.info(
f"[EPLB] finished update expert weight for layer: {self.layer_id}")
logger.debug(
"[eplb/d2d_loader] Layer %s D2D transfer completed, updated_experts=%s",
self.layer_id,
len(self.recv_expert_list),
)
if self.layer_id == self.eplb_adaptor.num_moe_layers - 1:
logger.info(
"[eplb/d2d_loader] Full expert weight update cycle completed, total_layers=%s",
self.eplb_adaptor.num_moe_layers,
)
self.recv_expert_list = []
self.updated_expert_map = None
self.layer_id = -1
self.state = ExpertWeightUpdateState.WAITING
def load_impl(self, old_expert_table, new_expert_table):
raise NotImplementedError

203
vllm_ascend/eplb/core/eplb_utils.py
View File

@@ -15,121 +15,122 @@
# This file is a part of the vllm-ascend project.
#
# Todo: Once https://github.com/vllm-project/vllm/issues/22246 is merged in vllm. Remove eplb utils.
import random
import json
from collections import defaultdict
import numpy as np
import torch
from vllm.logger import logger
from vllm.model_executor.layers.fused_moe.expert_map_manager import determine_expert_map
def determine_default_expert_map(global_expert_num, world_size, rank_id,
global_redundant_expert_num):
if world_size == 1:
local_ids = torch.arange(global_expert_num, dtype=torch.int32)
return (global_expert_num, local_ids)
def expert_file_to_tensor(expert_map_path, layer_id):
with open(expert_map_path) as f:
data = json.load(f)
physical_count = 0
device_data = []
if layer_id > data["moe_layer_count"]:
raise ValueError("Invalid EPLB Table")
if layer_id == data["moe_layer_count"]:
logger.warning("[eplb/utils] Init expert map of mtp/eagle when using sample.")
for device in data["layer_list"][0]["device_list"]:
physical_count += len(device["device_expert"])
return None, physical_count
for device in data["layer_list"][layer_id]["device_list"]:
physical_count += len(device["device_expert"])
device_data.append(device["device_expert"])
global_placement = torch.tensor(device_data, dtype=torch.int32)
return global_placement, physical_count
local_num_experts = global_expert_num // world_size
expert_map = torch.full((global_expert_num, ), -1, dtype=torch.int32)
if rank_id < world_size - 1:
start = rank_id * local_num_experts
end = (rank_id + 1) * local_num_experts
local_count = local_num_experts
else:
start = rank_id * local_num_experts
end = global_expert_num
local_count = global_expert_num - rank_id * local_num_experts
if isinstance(global_redundant_expert_num,
int) and rank_id < global_redundant_expert_num:
local_count += 1
if end < global_expert_num:
end += 1
def generate_global_placement(n_expert, ep_size, n_redundant, num_shared_experts):
n_expert -= num_shared_experts
if (n_expert + n_redundant) % ep_size != 0:
raise ValueError("(n_expert + n_redundant) % ep_size must be 0")
all_experts = np.arange(n_expert)
groups = np.array_split(all_experts, ep_size)
for i in range(n_redundant):
j = i % ep_size + 1
if len(groups[-j]) == 0:
groups[-j] = np.append(groups[-j], j)
else:
start -= 1
if isinstance(local_count, int):
local_ids = torch.arange(local_count, dtype=torch.int32)
expert_map[start:end] = local_ids
return (local_count, expert_map)
groups[-j] = np.append(groups[-j], (groups[-j][-1] + 1) % n_expert)
if num_shared_experts > 0:
for i, group in enumerate(groups):
groups[i] = np.append(group, n_expert + i % num_shared_experts)
return torch.tensor(groups, dtype=torch.int32)
def generate_log2phy_map(expert_map):
num_local_experts = expert_map.max() + 1
log2phy_map = expert_map.clone()
num_ranks, num_global_expert = log2phy_map.shape
def init_eplb_config(eplb_config, layer_id, moe_config, mix_placement=False, num_shared_experts=1, tp_size=None):
expert_map_path = eplb_config.expert_map_path
n_experts = moe_config.num_experts
ep_size = moe_config.ep_size
global_placement = None
eplb_enable = eplb_config.dynamic_eplb
n_redundant = eplb_config.num_redundant_experts if eplb_enable else 0
num_shared_experts = num_shared_experts if mix_placement else 0
row_indices = torch.arange(num_ranks).view(-1, 1).expand(num_ranks, \
num_global_expert) * num_local_experts
log2phy_map[log2phy_map != -1] += row_indices[log2phy_map != -1]
if ep_size == 1:
assert not eplb_enable, "EPLB must used in expert parallelism."
return None, None, None, n_redundant
for idx in range(num_global_expert):
positive_rank_idx = torch.where(log2phy_map[:, idx] != -1)[0]
negative_rank_idx = torch.where(log2phy_map[:, idx] == -1)[0]
num_rank_holding_expert = positive_rank_idx.size(0)
if expert_map_path:
eplb_enable = True
global_placement, physical_count = expert_file_to_tensor(expert_map_path, layer_id)
n_redundant = physical_count - n_experts
elif not eplb_enable:
_, expert_map, _ = determine_expert_map(ep_size, moe_config.ep_rank, n_experts)
return None, expert_map, None, 0
if num_rank_holding_expert == 0:
log2phy_map[:, idx] = torch.full((num_ranks, ),
0,
dtype=log2phy_map.dtype)
if global_placement is None:
global_placement = generate_global_placement(n_experts, ep_size, n_redundant, num_shared_experts)
if mix_placement:
n_redundant += ep_size - 1
global_expert_map = []
for rankid in range(ep_size):
expert_map = torch.full((n_experts,), -1, dtype=torch.int32)
local_placement = global_placement[rankid]
expert_map[local_placement] = torch.arange(local_placement.shape[0], dtype=torch.int32)
global_expert_map.append(expert_map)
if rankid == moe_config.ep_rank:
local_expert_map = expert_map
log2phy = (
generate_log2phy_map(
global_expert_map,
moe_config.ep_rank,
tp_size=int(tp_size) if tp_size is not None else None,
).npu()
if eplb_enable
else None
)
if num_rank_holding_expert == 1:
log2phy_map[negative_rank_idx, idx] = torch.full(
(num_ranks - 1, ),
log2phy_map[positive_rank_idx, idx].item(),
dtype=log2phy_map.dtype)
return torch.stack(global_expert_map), local_expert_map, log2phy, n_redundant
def generate_log2phy_map(global_expert_map, ep_rank, tp_size: int | None = None):
log2phy_map = defaultdict(list)
valid_count = torch.sum(global_expert_map[0] != -1)
for rankid, map_per_rank in enumerate(global_expert_map):
for idx, val in enumerate(map_per_rank):
val = val.item()
if val != -1:
log2phy_map[idx].append(val + rankid * valid_count)
for key in log2phy_map:
num_of_duplications = len(log2phy_map[key])
if tp_size is not None and tp_size > 1:
tp_rank = ep_rank % tp_size
dp_like_rank = ep_rank // tp_size
replica_index = (tp_rank + dp_like_rank + key) % num_of_duplications
else:
try:
random_list = [
random.choice(log2phy_map[positive_rank_idx, idx])
for _ in range(num_ranks - num_rank_holding_expert)
]
log2phy_map[negative_rank_idx,
idx] = torch.tensor(random_list,
dtype=log2phy_map.dtype)
except Exception as e:
logger.error(f"Fail to get log2phy_map: {str(e)}")
replica_index = ep_rank % num_of_duplications
log2phy_map[key] = log2phy_map[key][replica_index]
log2phy_map = torch.scatter(
torch.zeros(len(log2phy_map), dtype=torch.int32),
0,
torch.tensor(list(log2phy_map), dtype=torch.int64),
torch.tensor(list(log2phy_map.values()), dtype=torch.int32),
)
return log2phy_map
def determine_default_log2phy_map(global_expert_num, world_size, rank_id,
global_redundant_expert_num):
if world_size == 1:
local_ids = torch.arange(global_expert_num, dtype=torch.int32)
expert_map_all = local_ids.unsqueeze(0).expand(world_size, -1)
log2phy_map_all = generate_log2phy_map(expert_map_all)
return log2phy_map_all[rank_id]
local_num_experts = global_expert_num // world_size
expert_map_all = torch.full((world_size, global_expert_num),
-1,
dtype=torch.int32)
for r in range(world_size):
if r < world_size - 1:
start = r * local_num_experts
end = (r + 1) * local_num_experts
local_count = local_num_experts
else:
start = r * local_num_experts
end = global_expert_num
local_count = global_expert_num - r * local_num_experts
if isinstance(global_redundant_expert_num,
int) and rank_id < global_redundant_expert_num:
local_count += 1
if end < global_expert_num:
end += 1
else:
start -= 1
if isinstance(local_count, int):
local_ids = torch.arange(local_count, dtype=torch.int32)
expert_map_all[r, start:end] = local_ids
log2phy_map_all = generate_log2phy_map(expert_map_all)
return log2phy_map_all[rank_id]

362
vllm_ascend/eplb/core/eplb_worker.py
View File

@@ -17,27 +17,31 @@
from multiprocessing import Process, Queue
from typing import Any
import networkx as nx # type: ignore
import numpy as np
import torch
import torch.distributed as dist
from vllm.distributed import get_ep_group
from vllm.logger import logger
from vllm_ascend.eplb.core.eplb_utils import generate_log2phy_map
from vllm_ascend.eplb.core.policy.policy_factory import (DynamicConfig,
PolicyFactory)
from vllm_ascend.eplb.core.policy.policy_factory import PolicyFactory
class EplbWorker:
def __init__(self, shared_dict, policy_type, enable_d2d: bool = True):
def __init__(
self,
shared_dict,
policy_type,
enable_d2d: bool = True,
tp_size: int | None = None,
):
self.policy_type = policy_type
self.policy = PolicyFactory.generate_policy(policy_type,
DynamicConfig())
self.policy = PolicyFactory.generate_policy(policy_type)
self.shared_dict = shared_dict
self.old_expert_maps = None
self.enable_d2d = enable_d2d
self.rank_id = dist.get_rank()
self.tp_size = tp_size
self.rank_id = get_ep_group().rank_in_group
self.multi_stage = policy_type == 3
def do_update(self):
# put data in to queue
@@ -59,13 +63,41 @@ class EplbWorker:
# Get MOE load information
load_info = self.fetch_and_sum_load_info()
if load_info is None:
logger.debug("[eplb/worker] No moe_load data available yet, skipping this cycle")
return
# Get the updated expert table based on the workload information
old_placement = self.global2local(self.old_expert_maps,
self.num_local_experts)
_, _, new_placement = self.calculate_rebalance_experts(
load_info, old_placement)
old_placement = self.global2local(self.old_expert_maps, self.num_local_experts)
_, _, new_placement = self.calculate_rebalance_experts(load_info, old_placement)
if self.rank_id == 0:
if self.multi_stage:
hotness = self._calculate_hotness(old_placement, load_info.sum(0))
else:
hotness = self._calculate_hotness(old_placement, load_info)
# ms-service-metric begin: expose EPLB hotness details for metrics collection.
current_mean, current_max, current_imbalance_list = self._compute_imbalance(
old_placement, hotness, return_list=True
)
update_mean, update_max, update_imbalance_list = self._compute_imbalance(
new_placement, hotness, return_list=True
)
self.latest_expert_hotness = {
"current_mean": current_mean,
"current_max": current_max,
"update_mean": update_mean,
"update_max": update_max,
"current_imbalance_list": current_imbalance_list,
"update_imbalance_list": update_imbalance_list,
}
# ms-service-metric end.
logger.info(
"[eplb/worker] Expert hotness imbalance, current: mean=%.3f max=%.3f, updated: mean=%.3f max=%.3f",
current_mean,
current_max,
update_mean,
update_max,
)
if not torch.is_tensor(new_placement):
new_placement = torch.tensor(new_placement)
@@ -73,10 +105,9 @@ class EplbWorker:
new_expert_maps = self.local2global(new_placement)
self.update_expert_map(new_expert_maps)
update_info = self.compose_expert_update_info_greedy(
new_expert_maps, self.old_expert_maps)
update_info = self.compose_expert_update_info_greedy(new_expert_maps, self.old_expert_maps)
self.old_expert_maps = new_expert_maps
logger.info("EPLB Process compute complete")
logger.debug("[eplb/worker] EPLB Process compute complete")
packed_update_info = self.pack_update_info(update_info)
@@ -88,11 +119,8 @@ class EplbWorker:
for layer_id in range(num_layers):
# check if any logical expert is not placed on any rank
if torch.unique(new_placement[layer_id]).numel() < torch.unique(
old_placement[layer_id]).numel():
logger.error(
f"There exists expert not placed on any rank in layer {layer_id}"
)
if torch.unique(new_placement[layer_id]).numel() < torch.unique(old_placement[layer_id]).numel():
logger.error("[eplb/worker] There exists expert not placed on any rank in layer %s", layer_id)
new_placement[layer_id] = old_placement[layer_id]
continue
@@ -101,134 +129,30 @@ class EplbWorker:
old_placement_check = old_placement[layer_id][rank_id]
# check if same logical experts are placed on the same NPU
if new_placement_check.numel() != torch.unique(
new_placement_check).numel():
if new_placement_check.numel() != torch.unique(new_placement_check).numel():
logger.error(
f"Replicated experts are placed on the same NPU, expert placement on layer {layer_id}, rank {rank_id} is invalid"
"[eplb/worker] Replicated experts are placed on the same NPU; "
"expert placement on layer %s, rank %s is invalid",
layer_id,
rank_id,
)
new_placement[layer_id] = old_placement[layer_id]
break
# check if there is any experts movement inside one NPU
expert_not_move = torch.isin(new_placement_check,
old_placement_check)
if not torch.equal(new_placement_check[expert_not_move],
old_placement_check[expert_not_move]):
expert_not_move = torch.isin(new_placement_check, old_placement_check)
if not torch.equal(new_placement_check[expert_not_move], old_placement_check[expert_not_move]):
logger.error(
f"There exists expert movement inside NPU, expert placement on layer {layer_id}, rank {rank_id} is invalid"
"[eplb/worker] Expert movement inside NPU detected; "
"expert placement on layer %s, rank %s is invalid",
layer_id,
rank_id,
)
new_placement[layer_id] = old_placement[layer_id]
break
def compose_expert_update_info_bipartite(self, updated_expert_maps_org,
current_expert_maps_org):
# transform numpy array to torch tensor
updated_expert_maps = updated_expert_maps_org.clone()
current_expert_maps = current_expert_maps_org.clone()
updated_expert_maps = np.array(updated_expert_maps)
current_expert_maps = np.array(current_expert_maps)
num_layers = current_expert_maps.shape[0]
for layer_id in range(num_layers):
updated_expert_maps_this_layer = updated_expert_maps[layer_id]
current_expert_maps_this_layer = current_expert_maps[layer_id]
updated_expert_maps_this_layer_org = updated_expert_maps_org[
layer_id]
from typing import Any
expert_send_info_this_layer: dict[Any, Any] = {}
expert_recv_info_this_layer: dict[Any, Any] = {}
# Guard Clause: if there is no expert weight update, avoid subsequent processing
if (np.equal(updated_expert_maps_this_layer,
current_expert_maps_this_layer)).all():
yield (expert_send_info_this_layer,
expert_recv_info_this_layer,
updated_expert_maps_this_layer_org, layer_id)
# Parse expert_ids each rank needs to receive from other ranks
dst_rank_indices, experts_to_recv = np.where(
(current_expert_maps_this_layer == -1)
& (updated_expert_maps_this_layer != -1))
# record src ranks for potential transfer
src_ranks_set = dict()
for idx in range(len(dst_rank_indices)):
expert_id = experts_to_recv[idx].item()
if expert_id not in src_ranks_set:
src_ranks_set[expert_id] = np.where(
current_expert_maps_this_layer[:, expert_id] != -1)[0]
# loop until all experts are scheduled
while len(dst_rank_indices) > 0:
# construct bipartite graph
graph_expert_update: nx.Graph = nx.Graph()
for idx in range(len(dst_rank_indices)):
dst_rank_id = dst_rank_indices[idx].item()
expert_id = experts_to_recv[idx].item()
# add src ranks
src_rank_ids = src_ranks_set[expert_id]
graph_expert_update.add_nodes_from(src_rank_ids,
bipartite=0)
# add dest rank
graph_expert_update.add_node(str(dst_rank_id), bipartite=1)
# add edges
for src_rank_id in src_rank_ids:
graph_expert_update.add_edge(src_rank_id,
str(dst_rank_id))
# graph may not be connected
connected_components = list(
nx.connected_components(graph_expert_update))
all_matches = {}
# matching in this loop
for i, component in enumerate(connected_components):
subgraph = graph_expert_update.subgraph(component)
component_matching = nx.bipartite.maximum_matching(
subgraph)
all_matches.update(component_matching)
for src_rank, dst_rank in all_matches.items():
dst_rank = int(dst_rank)
assert src_rank != dst_rank
if graph_expert_update.nodes[src_rank]['bipartite'] == 0:
# currently not scheduled experts in rank dst_rank
experts_v = experts_to_recv[np.where(
dst_rank_indices == dst_rank)]
# src: src_rank, dest: dst_rank, expert: expert_id
expert_id = np.intersect1d(
experts_v,
np.where(current_expert_maps_this_layer[src_rank]
!= -1))[0]
# record send/rcv pairs
if src_rank not in expert_send_info_this_layer:
expert_send_info_this_layer[src_rank] = []
if dst_rank not in expert_recv_info_this_layer:
expert_recv_info_this_layer[dst_rank] = []
expert_send_info_this_layer[src_rank].append(
(dst_rank, expert_id))
expert_recv_info_this_layer[dst_rank].append(
(src_rank, expert_id))
remove_index = np.where(
np.logical_and(dst_rank_indices == dst_rank,
experts_to_recv == expert_id))
# update
dst_rank_indices = np.delete(dst_rank_indices,
remove_index)
experts_to_recv = np.delete(experts_to_recv,
remove_index)
yield (expert_send_info_this_layer, expert_recv_info_this_layer,
updated_expert_maps_this_layer_org, layer_id)
# TODO: Here only expert weight exchange is considered, need to be extended to cover other weight update cases
def compose_expert_update_info_greedy(self, updated_expert_maps,
current_expert_maps):
def compose_expert_update_info_greedy(self, updated_expert_maps, current_expert_maps):
num_layers = current_expert_maps.shape[0]
for layer_id in range(num_layers):
updated_expert_maps_this_layer = updated_expert_maps[layer_id]
@@ -238,19 +162,24 @@ class EplbWorker:
expert_recv_info_this_layer: dict[Any, Any] = {}
# Guard Clause: if there is no expert weight update, avoid subsequent processing
if torch.equal(updated_expert_maps_this_layer,
current_expert_maps_this_layer):
yield (expert_send_info_this_layer,
expert_recv_info_this_layer,
updated_expert_maps_this_layer, layer_id)
if torch.equal(updated_expert_maps_this_layer, current_expert_maps_this_layer):
yield (
expert_send_info_this_layer,
expert_recv_info_this_layer,
updated_expert_maps_this_layer,
layer_id,
)
continue
# Parse expert_ids each rank needs to receive from other ranks
dst_rank_indices, experts_to_recv = torch.where((current_expert_maps_this_layer == -1) \
& (updated_expert_maps_this_layer != -1))
dst_rank_indices, experts_to_recv = torch.where(
(current_expert_maps_this_layer == -1) & (updated_expert_maps_this_layer != -1)
)
# Parse expert_ids each rank needs to send to other ranks
src_rank_indices, experts_to_send = torch.where((current_expert_maps_this_layer != -1) \
& (updated_expert_maps_this_layer == -1))
src_rank_indices, experts_to_send = torch.where(
(current_expert_maps_this_layer != -1) & (updated_expert_maps_this_layer == -1)
)
for idx in range(len(dst_rank_indices)):
dst_rank_id = dst_rank_indices[idx].item()
@@ -258,27 +187,27 @@ class EplbWorker:
if dst_rank_id not in expert_recv_info_this_layer:
expert_recv_info_this_layer[dst_rank_id] = []
if not torch.isin(torch.tensor(expert_id),
experts_to_send).any():
if not torch.isin(torch.tensor(expert_id), experts_to_send).any():
# if expert_id are not sent out from any npu, it will be copied from one npu holding this expert
candidate_src_rank_indices = torch.where(
current_expert_maps_this_layer[:, expert_id] != -1)[0]
candidate_src_rank_indices = torch.where(current_expert_maps_this_layer[:, expert_id] != -1)[0]
else:
candidate_src_rank_indices = src_rank_indices[
experts_to_send == expert_id]
candidate_src_rank_indices = src_rank_indices[experts_to_send == expert_id]
# TODO: improve selection criterion of npu sending expert_id considering such as intra-node or inter-node...
# TODO: improve selection criterion of NPU sending expert_id,
# considering intra-node or inter-node...
src_rank_id = candidate_src_rank_indices[0].item()
if src_rank_id not in expert_send_info_this_layer:
expert_send_info_this_layer[src_rank_id] = []
expert_send_info_this_layer[src_rank_id].append(
(dst_rank_id, expert_id))
expert_recv_info_this_layer[dst_rank_id].append(
(src_rank_id, expert_id))
expert_send_info_this_layer[src_rank_id].append((dst_rank_id, expert_id))
expert_recv_info_this_layer[dst_rank_id].append((src_rank_id, expert_id))
yield (expert_send_info_this_layer, expert_recv_info_this_layer,
updated_expert_maps_this_layer, layer_id)
yield (
expert_send_info_this_layer,
expert_recv_info_this_layer,
updated_expert_maps_this_layer,
layer_id,
)
def calculate_rebalance_experts(self, load_info, old_placement):
"""
@@ -287,8 +216,7 @@ class EplbWorker:
if self.old_expert_maps is None:
return False, None, None
changed, priority, new_map = self.policy.rebalance_experts(
old_placement, load_info)
changed, priority, new_map = self.policy.rebalance_experts(old_placement, load_info)
return changed, priority, new_map
def get_init_expert_maps(self):
@@ -305,19 +233,13 @@ class EplbWorker:
return self.shared_dict.get("moe_load", None)
def update_expert_map(self, expert_maps):
self.shared_dict["expert_maps"] = expert_maps
def global2local(self, placement: torch.Tensor,
E_local: int) -> tuple[torch.Tensor, torch.Tensor]:
def global2local(self, placement: torch.Tensor, E_local: int) -> tuple[torch.Tensor, torch.Tensor]:
L, G, _ = placement.shape
device = placement.device
pt_local = torch.full((L, G, E_local),
fill_value=-1,
dtype=torch.long,
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)
@@ -329,7 +251,6 @@ class EplbWorker:
return pt_local
def local2global(self, placement_local: torch.Tensor) -> torch.Tensor:
L, G, E_local = placement_local.shape
device = placement_local.device
@@ -339,10 +260,7 @@ class EplbWorker:
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)
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)
@@ -363,30 +281,67 @@ class EplbWorker:
layer_ids = []
for send_info, recv_info, new_expert_map, layer_id in update_info_generator:
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_info_this_rank = send_info.get(self.rank_id, [])
recv_info_this_rank = recv_info.get(self.rank_id, [])
send_all.append(send_info_this_rank)
recv_all.append(recv_info_this_rank)
maps.append(new_expert_map[self.rank_id].numpy().tolist())
log2phy_map = generate_log2phy_map(new_expert_map)
log2phy_all.append(log2phy_map[self.rank_id].numpy().tolist())
log2phy_map = generate_log2phy_map(
new_expert_map,
self.rank_id,
tp_size=self.tp_size,
)
log2phy_all.append(log2phy_map.numpy().tolist())
layer_ids.append(layer_id)
return list(zip(send_all, recv_all, maps, log2phy_all, layer_ids))
@staticmethod
def _compute_imbalance(deployment_all_layer, hotness_all_layer: np.ndarray, return_list: bool = False):
imbalance_list = []
deployment_all_layer = np.array(deployment_all_layer)
for deployment, hotness in zip(deployment_all_layer, hotness_all_layer):
counts = np.bincount(deployment.reshape(-1), minlength=hotness.shape[0])
unit_hotness = np.divide(hotness, counts, out=np.zeros_like(hotness, dtype=float), where=counts != 0)
stage_load = unit_hotness[deployment].sum(-1)
stage_par = stage_load.max() / stage_load.mean()
imbalance_list.append(stage_par)
max_val = max(imbalance_list)
mean_val = sum(imbalance_list) / len(imbalance_list)
# ms-service-metric begin: optionally expose per-layer imbalance without recomputing it.
if return_list:
return mean_val, max_val, imbalance_list
# ms-service-metric end.
return mean_val, max_val
@staticmethod
def _calculate_hotness(deployment_all_layer, moe_load_all_layer):
hotnesses = []
num_of_expert = deployment_all_layer.shape[1] * deployment_all_layer.shape[2]
for deployment, rank_load in zip(deployment_all_layer, moe_load_all_layer.numpy()):
hotness = np.zeros(num_of_expert, dtype=rank_load.dtype)
deployment_flat = deployment.ravel()
rank_load_flat = rank_load.ravel()
np.add.at(hotness, deployment_flat, rank_load_flat)
hotnesses.append(hotness)
return np.array(hotnesses)
class EplbProcess:
def __init__(self,
shared_dict,
policy_type: int = 0,
enable_d2d: bool = True):
def __init__(
self,
shared_dict,
policy_type: int = 0,
enable_d2d: bool = True,
tp_size: int | None = None,
):
"""
Args:
shared_dict: Cross-process shared dict returned by Manager().dict()
@@ -400,13 +355,31 @@ class EplbProcess:
self.block_update_q: Queue[Any] = Queue(maxsize=1)
# Create EplbWorker instance
self.worker = EplbWorker(self.shared_dict, self.policy_type,
self.enable_d2d)
self.worker = EplbWorker(
self.shared_dict,
self.policy_type,
self.enable_d2d,
tp_size=tp_size,
)
def worker_process(self, planner_q, block_update_q):
"""
Subprocess entry: bind to specified NPU, loop waiting for planner_q to wake up, call do_update, then notify main process update is complete.
Subprocess entry: bind to specified NPU, loop waiting for planner_q to wake up,
call do_update, then notify main process update is complete.
"""
try:
from ms_service_metric.adapters.vllm.adapter import get_vllm_adapter, initialize_vllm_metric # type: ignore
initialize_vllm_metric()
adapter = get_vllm_adapter()
logger.info("[EPLB metrics] The adapter initialized: %s", adapter.is_initialized())
except Exception as e:
logger.warning("[EPLB metrics] Failed to initialize metrics: %s", e)
if self.policy_type == 3:
from vllm_ascend.eplb.core.policy.policy_flashlb import warm_up
warm_up()
while True:
try:
planner_q.get()
@@ -420,17 +393,18 @@ class EplbProcess:
break
except Exception as e:
logger.warning(f"[EPLB subprocess Exiting due to error: {e}",
exc_info=True)
logger.warning(
"[eplb/worker] Subprocess crashed, EPLB optimization will stop. error=%s",
e,
exc_info=True,
)
break
def _launch_process(self):
"""
Use spawn method to launch subprocess and return (planner_q, block_update_q, proc).
"""
proc = Process(target=self.worker_process,
args=(self.planner_q, self.block_update_q),
daemon=True)
proc = Process(target=self.worker_process, args=(self.planner_q, self.block_update_q), daemon=True)
proc.start()
return proc

12
vllm_ascend/eplb/core/policy/policy_abstract.py
View File

@@ -3,19 +3,7 @@
from abc import abstractmethod
class DynamicConfig:
placement_policy = None
max_transferred_expert_per_layer = 100 # Maximum number of experts that can be migrated per layer on a single host
ep_worldsize = 64 # Total number of dies across the entire cluster where experts are distributed
num_die_per_host = 8 # Number of dies on each host machine
class EplbPolicy:
def __init__(self, config: DynamicConfig):
self.config = config
@abstractmethod
def rebalance_experts(self, current_expert_table, expert_workload):
"""

350
vllm_ascend/eplb/core/policy/policy_default_eplb.py Normal file
View File

@@ -0,0 +1,350 @@
# Copyright Huawei Technologies Co., Ltd. 2024-2025. All rights reserved.
# Todo: Once https://github.com/vllm-project/vllm/pull/24069 is merged in vllm. Remove this policy.
from collections import defaultdict
from typing import cast
import numpy as np
from .policy_abstract import EplbPolicy
class DynamicTable:
# workload_table:
# 3D matrix: [layer, gpus, experts_per_gpu_per_layer] -> value: workload (heat) at the corresponding position
# Size: number of layers * number of GPUs * number of experts per GPU per layer
# The element at (i, j, k) represents the workload (heat) of the k-th expert on the j-th GPU in the i-th layer
# For experts that are not available or collected, the value is set to -1
workload_table = None
# placement_table:
# 3D matrix: [layer, gpus, experts_per_gpu_per_layer] -> value: physical expert ID at the corresponding position
# Size: number of layers * number of GPUs * number of experts per GPU per layer
# The element at (i, j, k) represents the physical expert ID of the k-th expert on the j-th GPU in the i-th layer
# For experts that are not available or collected, the value is set to -1
placement_table = None
class DefaultEplb(EplbPolicy):
@staticmethod
def add_redundant(current_expert_table, expert_workload, num_original_expert):
layer_num, npu_num, experts_per_npu = expert_workload.shape
workload_new = np.zeros((layer_num, num_original_expert))
for layer_idx in range(layer_num):
workload_dict: dict[int, int] = defaultdict(int)
placement_layer = current_expert_table[layer_idx].copy()
workload_layer = expert_workload[layer_idx].copy()
for npu_idx in range(npu_num):
for expert_idx in range(experts_per_npu):
workload_dict[placement_layer[npu_idx][expert_idx]] += workload_layer[npu_idx][expert_idx]
for expert_idx in range(num_original_expert):
workload_new[layer_idx][expert_idx] = workload_dict[expert_idx]
return workload_new
@staticmethod
# Split hot (high-load) experts into redundant experts
def original_compute_balanced_pack_redundancy(origin_weights, card_num, num_redundancy_expert):
# Step 1: Sort the items by weight in descending order (we are sorting by weight now)
# Sort based on the second element (the second value of each tuple)
route_expert_num = len(origin_weights)
route_expert_redundancy: list[list[int]] = [[] for _ in range(route_expert_num)]
for i in range(num_redundancy_expert):
sorted_indices = np.argsort([t[1] for t in origin_weights], kind="stable")[::-1]
weights = [origin_weights[idx] for idx in sorted_indices]
tmp_raw_weight = weights[0][1] * (len(route_expert_redundancy[weights[0][0]]) + 1)
route_expert_redundancy[weights[0][0]].append(route_expert_num + i)
avg_weight = tmp_raw_weight / (len(route_expert_redundancy[weights[0][0]]) + 1)
weights[0] = (weights[0][0], avg_weight)
origin_weights = weights
# Step 2: Calculate the number of items per box
expert_num = route_expert_num + num_redundancy_expert
items_per_box = expert_num // card_num # Number of items per box
remaining_items = expert_num % card_num # Number of items per box
# Step 3: Initialize card_num boxes with empty lists to store item IDs
boxes: list[list[int]] = [[] for _ in range(card_num)]
boxes_weights: list[list[float]] = [[] for _ in range(card_num)]
box_weights = [0] * card_num # To store the total weight of each box
box_counts = [0] * card_num # To store the number of items in each box
index = 0
for i in range(route_expert_num):
redundancy_num = len(route_expert_redundancy[i])
for _ in range(redundancy_num):
cur_weight = 0
for item, weight in origin_weights:
if item == i:
cur_weight = weight
boxes[index].append(i)
boxes_weights[index].append(cur_weight)
box_weights[index] += cur_weight
box_counts[index] += 1
index += 1
sorted_indices = np.argsort([t[1] for t in origin_weights], kind="stable")[::-1]
origin_weights = [origin_weights[idx] for idx in sorted_indices]
# Step 4: Distribute items into boxes based on weight
for item_id, weight in origin_weights:
# Find the box with the least items but not full
min_box_index = -1
for i in range(card_num):
if item_id in boxes[i]:
continue
# Only choose boxes that still have space (box_counts[i] < items_per_box)
if box_counts[i] < items_per_box or (box_counts[i] == items_per_box and remaining_items > 0):
if min_box_index == -1 or box_weights[i] < box_weights[min_box_index]:
min_box_index = i
# Place the item (id) into the selected box
boxes[min_box_index].append(item_id)
boxes_weights[min_box_index].append(weight)
box_weights[min_box_index] += weight
box_counts[min_box_index] += 1
# If there's an imbalance in the remaining items, reduce the "remaining_items" counter
if box_counts[min_box_index] == (items_per_box + 1) and remaining_items > 0:
remaining_items -= 1
# Step 5: Output each box's contents and total weight
result = []
for i in range(card_num):
result.append(
{
"box_index": i + 1,
"items": boxes[i], # List of item IDs in the box
"weight": boxes_weights[i],
"total_weight": box_weights[i], # Total weight in this box
"item_count": box_counts[i], # Number of items in the box
}
)
return result, boxes
# Split hot (high-load) experts into redundant experts
@staticmethod
def compute_balanced_pack_redundancy(origin_weights, card_num, num_redundancy_expert):
route_expert_num = len(origin_weights)
route_expert_redundancy: list[list[int]] = [[] for _ in range(route_expert_num)]
for i in range(num_redundancy_expert):
sorted_indices = np.argsort([t[1] for t in origin_weights], kind="stable")[::-1]
weights = [origin_weights[idx] for idx in sorted_indices]
tmp_raw_weight = weights[0][1] * (len(route_expert_redundancy[weights[0][0]]) + 1)
route_expert_redundancy[weights[0][0]].append(route_expert_num + i)
avg_weight = tmp_raw_weight / (len(route_expert_redundancy[weights[0][0]]) + 1)
weights[0] = (weights[0][0], avg_weight)
origin_weights = weights
expert_num = route_expert_num + num_redundancy_expert
if card_num == 0:
raise RuntimeError("card_num can not be 0.")
items_per_box = expert_num // card_num
remaining_items = expert_num % card_num
boxes: list[list[int]] = [[] for _ in range(card_num)]
boxes_weights: list[list[float]] = [[] for _ in range(card_num)]
box_weights = [0] * card_num
box_counts = [0] * card_num
all_weights = np.zeros((expert_num,), dtype="object")
all_weights[:route_expert_num] = origin_weights
index = route_expert_num
for i in range(route_expert_num):
redundancy_num = len(route_expert_redundancy[i])
for _ in range(redundancy_num):
for item, weight in origin_weights:
if item == i:
all_weights[index] = (item, weight)
index += 1
sorted_indices = np.argsort([t[1] for t in all_weights], kind="stable")[::-1]
all_weights = [all_weights[idx] for idx in sorted_indices]
for item_id, weight in all_weights:
min_box_index = -1
for i in range(card_num):
if box_counts[i] < items_per_box or (box_counts[i] == items_per_box and remaining_items > 0):
if min_box_index == -1 or box_weights[i] < box_weights[min_box_index]:
if item_id not in boxes[i]:
min_box_index = i
boxes[min_box_index].append(item_id)
boxes_weights[min_box_index].append(weight)
box_weights[min_box_index] += weight
box_counts[min_box_index] += 1
if box_counts[min_box_index] == (items_per_box + 1) and remaining_items > 0:
remaining_items -= 1
result = []
for i in range(card_num):
result.append(
{
"box_index": i + 1,
"items": boxes[i],
"weight": boxes_weights[i],
"total_weight": box_weights[i],
"item_count": box_counts[i],
}
)
return result, boxes
# Scheme without redundant experts
@staticmethod
def compute_balanced_pack(origin_weights, card_num):
sorted_indices = np.argsort([t[1] for t in origin_weights])[::-1]
weights = origin_weights[sorted_indices]
expert_num = len(weights)
if card_num == 0:
raise RuntimeError("card_num can not be 0.")
items_per_box = expert_num // card_num
remaining_items = expert_num % card_num
boxes: list[list[int]] = [[] for _ in range(card_num)]
boxes_weights: list[list[float]] = [[] for _ in range(card_num)]
box_weights = [0] * card_num
box_counts = [0] * card_num
for item_id, weight in weights:
min_box_index = -1
for i in range(card_num):
if box_counts[i] < items_per_box or (box_counts[i] == items_per_box and remaining_items > 0):
if min_box_index == -1 or box_weights[i] < box_weights[min_box_index]:
min_box_index = i
boxes[min_box_index].append(item_id)
boxes_weights[min_box_index].append(weight)
box_weights[min_box_index] += weight
box_counts[min_box_index] += 1
if box_counts[min_box_index] == (items_per_box + 1) and remaining_items > 0:
remaining_items -= 1
result = []
for i in range(card_num):
result.append(
{
"box_index": i + 1,
"items": boxes[i],
"weight": boxes_weights[i],
"total_weight": box_weights[i],
"item_count": box_counts[i],
}
)
return result, boxes
@staticmethod
def get_redundant_num(npu_num, counts):
redundant_num_each_npu: int = np.sum(counts - 1)
return redundant_num_each_npu
@staticmethod
def calculate_max_heat_per_layer(workload_table, layer_num):
max_heat_per_layer: list[float] = []
for layer_idx in range(layer_num):
npu_heats_now = np.sum(workload_table[layer_idx], axis=1)
max_heat_per_layer.append(np.max(npu_heats_now))
return max_heat_per_layer
@staticmethod
def constraint_expert_local_exchange(current_expert_table, global_deployment):
for layer_id in range(len(global_deployment)):
for card_id in range(len(global_deployment[layer_id])):
current_list = [int(x) for x in current_expert_table[layer_id][card_id]]
new_list = [int(x) for x in global_deployment[layer_id][card_id]]
num = len(new_list)
new_index = [-1] * num
new_result = [-1] * num
remaining_elements = []
for i in range(num):
flag = True
for j in range(num):
if new_list[i] == current_list[j] and new_index[j] == -1:
new_index[j] = 0
new_result[j] = current_list[j]
flag = False
break
if flag:
remaining_elements.append(new_list[i])
index = 0
for k in range(num):
if new_result[k] == -1:
new_result[k] = remaining_elements[index]
index += 1
global_deployment[layer_id][card_id] = new_result
return global_deployment
def rebalance_experts(self, current_expert_table, expert_workload):
info = DynamicTable()
info.workload_table = np.array(expert_workload)
info.placement_table = np.array(current_expert_table)
assert info.workload_table is not None
layer_num, num_npus, experts_per_npu = info.workload_table.shape
assert info.placement_table is not None
row = cast(np.ndarray, info.placement_table[0])
expert_ids, counts = np.unique(row, return_counts=True)
num_redundancy_expert = self.get_redundant_num(num_npus, counts)
num_original_expert = len(expert_ids)
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)
# Perform load balancing and deploy redundant experts
layer_num = layer_workloads.shape[0]
expert_num = layer_workloads.shape[1]
# Validate that the number of experts, number of cards, and number of redundant experts
# do not exceed the number of cards.
if num_original_expert != expert_num:
raise ValueError(
f"the number of original experts {num_original_expert} must be equal to expert_num {expert_num}"
)
if num_npus <= 0:
raise ValueError("the number of NPUs must be greater than 0")
if num_npus < num_redundancy_expert:
raise ValueError(
"the number of NPUs "
f"{num_npus} must be greater than or equal to the number of redundant experts "
f"{num_redundancy_expert}"
)
# Number of experts deployed on each card includes one redundant expert
global_deployment: list[list[list[int]]] = [[[] for _ in range(num_npus)] for _ in range(layer_num)]
# Iterate to obtain the placement strategy for each layer, taking computational balance into account
max_heat_per_layer_after = np.zeros([layer_num])
for layer in range(layer_num):
# Get the expert IDs and their corresponding workloads for the current layer;
# workloads need to be normalized, and one redundant expert is added per card
weights = np.zeros((expert_num,), dtype="object")
for expert_id, workload_weight in enumerate(layer_workloads[layer]):
weights[expert_id] = (expert_id, workload_weight)
# Obtain the globally balanced placement strategy for each layer
result, layer_deployment = self.original_compute_balanced_pack_redundancy(
weights, num_npus, num_redundancy_expert
)
global_deployment[layer] = layer_deployment
max_heat_per_layer_after[layer] = max(result, key=lambda x: x["total_weight"])["total_weight"]
new_global_deployment = self.constraint_expert_local_exchange(current_expert_table, global_deployment)
# Obtain the priority of each layer
layer_changed_ratio = []
for layer_idx in range(layer_num):
layer_changed_ratio.append(max_heat_per_layer_after[layer_idx] / max_heat_per_layer_before[layer_idx])
per_layer_priority = np.argsort(layer_changed_ratio)
npu_heat_all_after = sum(max_heat_per_layer_after)
change = 0
if npu_heat_all_after < 0.95 * npu_heat_all_origin:
change = 1
return change, per_layer_priority, np.array(new_global_deployment).tolist()

48
vllm_ascend/eplb/core/policy/policy_factory.py
View File

@@ -1,33 +1,41 @@
# Copyright Huawei Technologies Co., Ltd. 2023-2024. All rights reserved.
# Todo: Once https://github.com/vllm-project/vllm/pull/24069 is merged in vllm. Remove this factory.
from .policy_abstract import DynamicConfig, EplbPolicy
from .policy_dynamic_ep import DynamicEplb
from .policy_dynamic_ep_v2 import DynamicEplbV2
from .policy_flashlb import FlashLB
from vllm.logger import logger
from .policy_abstract import EplbPolicy
from .policy_default_eplb import DefaultEplb
from .policy_flashlb import FlashLB, warm_up
from .policy_random import RandomLoadBalance
from .policy_swift_balancer import SwiftBalanceEplb
class PolicyFactory:
@staticmethod
def generate_policy(policy_type: int, config: DynamicConfig) -> EplbPolicy:
policy = {
def generate_policy(policy_type: int) -> EplbPolicy:
policy: dict[int, type[EplbPolicy]] = {
# Constraint applying Dynamic EPLB policy V2:
# If there exists redundant expert:
# only one redundant expert can be placed in one NPU and its physical expert index must be 0
# Applying greedy d2d expert weight update composing
0:
RandomLoadBalance, # RandomLoadBalance: shuffle last physical expert on NPU 1 and 3
1:
DynamicEplb, # Dynamic EPLB policy: overall expert replacement based on current moe load
2:
DynamicEplbV2, # Dynamic EPLB policy V2: expert replacement with constrained number of expert shuffle
3:
FlashLB, # FlashLB EPLB policy: expert replacement based on Joint Optimization, Multi-Shot Enhancement and Incremental Adjustment
0: RandomLoadBalance, # RandomLoadBalance: shuffle last physical expert on NPU 1 and 3
1: DefaultEplb, # Dynamic EPLB policy: overall expert replacement based on current moe load
# Dynamic EPLB policy V2: expert replacement with constrained number of expert shuffle
2: SwiftBalanceEplb,
# FlashLB EPLB policy: expert replacement based on Joint Optimization,
# Multi-Shot Enhancement and Incremental Adjustment
3: FlashLB,
}
policy_class = policy.get(policy_type, RandomLoadBalance)
policy_instance = policy_class(config)
policy_class = policy.get(policy_type)
if policy_class is None:
policy_class = RandomLoadBalance
logger.warning(
"[eplb/policy] Unrecognized policy_type=%s, falling back to %s",
policy_type,
policy_class.__name__,
)
else:
logger.info("[eplb/policy] Policy: %s (type=%s)", policy_class.__name__, policy_type)
policy_instance = policy_class()
if policy_type == 3:
policy_instance.warm_up()
return policy_instance
warm_up()
return policy_instance

1501
vllm_ascend/eplb/core/policy/policy_flashlb.py
View File

File diff suppressed because it is too large Load Diff

6
vllm_ascend/eplb/core/policy/policy_random.py
View File

@@ -3,16 +3,12 @@
import copy
import random
from .policy_abstract import DynamicConfig, EplbPolicy
from .policy_abstract import EplbPolicy
random.seed(42)
class RandomLoadBalance(EplbPolicy):
def __init__(self, config: DynamicConfig):
super().__init__(config)
def rebalance_experts(self, current_expert_table, expert_workload):
new_table = copy.deepcopy(current_expert_table)
num_layers = len(current_expert_table)

751
vllm_ascend/eplb/core/policy/policy_swift_balancer.py Normal file
View File

@@ -0,0 +1,751 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections import defaultdict
import numpy as np
import torch
import torch_npu # noqa: F401
from .policy_abstract import EplbPolicy
class DynamicTable:
# workload_table:
# 3D matrix: [layer, gpus, experts_per_gpu_per_layer] -> value: workload (heat) at the corresponding position
# Size: number of layers * number of GPUs * number of experts per GPU per layer
# The element at (i, j, k) represents the workload (heat) of the k-th expert on the j-th GPU in the i-th layer
# For experts that are not available or collected, the value is set to -1
workload_table = None
# placement_table:
# 3D matrix: [layer, gpus, experts_per_gpu_per_layer] -> value: physical expert ID at the corresponding position
# Size: number of layers * number of GPUs * number of experts per GPU per layer
# The element at (i, j, k) represents the physical expert ID of the k-th expert on the j-th GPU in the i-th layer
# For experts that are not available or collected, the value is set to -1
placement_table = None
class SwiftBalanceEplb(EplbPolicy):
def __init__(self):
self.num_layers: int = 0
self.num_original_experts: int = 0
self.num_ranks: int = 0
self.num_experts_per_rank: int = 0
self.num_nodes: int = 0
self.is_node_redundant: bool = False
self.num_max_com: int = 1
self.imbalance_threshold: float = 1.01
self.increment = 0.01
self.swap_threshold: float = 0
self.max_swap_times: int = 100
self.num_die_per_host = torch.npu.device_count()
@staticmethod
def calculate_max_heat_per_layer(workload_table: np.ndarray) -> list[float]:
max_heat_per_layer: list[float] = []
for layer_idx in range(workload_table.shape[0]):
npu_heats_now = np.sum(workload_table[layer_idx], axis=1)
max_heat_per_layer.append(np.max(npu_heats_now))
return max_heat_per_layer
@staticmethod
def get_original_workload(
current_expert_table: np.ndarray, expert_workload: np.ndarray, num_original_expert: int
) -> np.ndarray:
"""
Accumulate workload for each routed expert
"""
layer_num, npu_num, experts_per_npu = expert_workload.shape
workload_new = np.zeros((layer_num, num_original_expert))
for layer_idx in range(layer_num):
workload_dict: dict[int, int] = defaultdict(int)
placement_layer = current_expert_table[layer_idx].copy()
workload_layer = expert_workload[layer_idx].copy()
for npu_idx in range(npu_num):
for expert_idx in range(experts_per_npu):
workload_dict[int(placement_layer[npu_idx][expert_idx])] += workload_layer[npu_idx][expert_idx]
for expert_idx in range(num_original_expert):
workload_new[layer_idx][expert_idx] = workload_dict[expert_idx]
return workload_new
@staticmethod
def constraint_expert_local_exchange(old_deployment: np.ndarray, new_deployment: np.ndarray):
"""
Align the new deployment with the old deployment
"""
for layer_id in range(len(new_deployment)):
for card_id in range(len(new_deployment[layer_id])):
current_list = [int(x) for x in old_deployment[layer_id][card_id]]
new_list = [int(x) for x in new_deployment[layer_id][card_id]]
num = len(new_list)
new_index = [-1] * num
new_result = [-1] * num
remaining_elements = []
for i in range(num):
flag = True
for j in range(num):
if new_list[i] == current_list[j] and new_index[j] == -1:
new_index[j] = 0
new_result[j] = current_list[j]
flag = False
break
if flag:
remaining_elements.append(new_list[i])
index = 0
for k in range(num):
if new_result[k] == -1:
new_result[k] = remaining_elements[index]
index += 1
new_deployment[layer_id][card_id] = new_result
def calculate_imbalance(self, cur_deployment: np.ndarray, cur_experts_load: np.ndarray) -> list[float]:
"""
Calculate the imbalance degree of each layer.
"""
per_layer_imbalance = []
num_per_expert = np.zeros_like(cur_experts_load)
for layer_id, layer in enumerate(cur_deployment):
for rank in layer:
for expert_id in rank:
num_per_expert[layer_id][expert_id] += 1
for layer_id, layer in enumerate(cur_deployment):
cur_layer_max_load = 0
total_load = 0
for rank in layer:
rank_load = 0
for expert_id in rank:
update_workload = cur_experts_load[layer_id][expert_id] / num_per_expert[layer_id][expert_id]
rank_load += update_workload
total_load += update_workload
if cur_layer_max_load < rank_load:
cur_layer_max_load = rank_load
avg_load = total_load / self.num_ranks
if abs(avg_load) < 1e-9:
cur_layer_imbalance = 1.0
else:
cur_layer_imbalance = cur_layer_max_load / avg_load
per_layer_imbalance.append(cur_layer_imbalance)
return per_layer_imbalance
def statistics_expert_distribution(
self, single_layer_deployment: np.ndarray
) -> tuple[list[list[int]], np.ndarray, set[int], int]:
"""
Statistics on the distribution of redundant experts and logical
experts under the current deployment
Parameters:
single_layer_deployment: [num_ranks, num_experts_per_rank]
the expert deployment status on each rank
Returns:
redundant_expert_pos: the positions of redundant experts
on each rank
expert_from_rank: [num_logical_experts] the rank where
each logical expert resides
num_redundant_experts: the number of redundant experts
"""
num_ranks = len(single_layer_deployment)
redundant_expert_pos: list[list[int]] = [[] for _ in range(num_ranks)]
num_redundant_experts = 0
expert_from_rank = np.zeros(self.num_original_experts, dtype=np.int64)
existing_experts = set()
for index in range(self.num_experts_per_rank):
for rank_id in range(num_ranks):
expert_id = int(single_layer_deployment[rank_id][index])
if expert_id not in existing_experts:
existing_experts.add(expert_id)
expert_from_rank[expert_id] = rank_id
else:
redundant_expert_pos[rank_id].append(index)
num_redundant_experts += 1
return (redundant_expert_pos, expert_from_rank, existing_experts, num_redundant_experts)
def compute_redundant_assignments(
self,
initial_weights: list,
num_redundant_experts: int,
num_ranks: int,
) -> tuple[list[tuple[int, float]], np.ndarray]:
"""
Reconfigure redundant experts based on current expert workload and
count the new expert workload after redundancy reconfiguration
Parameters:
initial_weights: [num_logical_experts] expert load statistics
num_redundant_experts: Number of redundant experts
num_ranks: Number of cards in the current node
Returns:
redundant_expert_list:[(expert_id, expert_load)]
the redundantly generated experts
update_weight: [num_logical_experts] expert workload status after
reconfiguring redundant experts
"""
current_weights = initial_weights.copy()
redundant_assignments = np.zeros(self.num_original_experts, dtype=np.int64)
for i in range(num_redundant_experts):
sorted_indices = np.argsort([w for _, w in current_weights], kind="stable")[::-1]
for index in sorted_indices:
target_expert = current_weights[index]
expert_id, original_weight = target_expert
current_redundancy = redundant_assignments[expert_id] + 1
if current_redundancy < num_ranks:
new_avg_weight = original_weight * (current_redundancy + 1) / (current_redundancy + 2)
redundant_assignments[expert_id] += 1
current_weights[index] = (expert_id, new_avg_weight)
break
update_weight = np.zeros(self.num_original_experts, dtype=np.float32)
for expert_id, expert_weight in current_weights:
update_weight[expert_id] = expert_weight
redundant_expert_list = []
if num_redundant_experts > 0:
for expert_id in range(self.num_original_experts):
for _ in range(redundant_assignments[expert_id]):
redundant_expert_list.append((expert_id, float(update_weight[expert_id])))
redundant_expert_list.sort(key=lambda x: x[1], reverse=True)
return redundant_expert_list, update_weight
def fill_in_undeployed_ranks(
self,
initial_weights: list,
rank_assignments: np.ndarray,
undeployed_ranks: list[int],
redundant_expert_pos: list[list[int]],
num_com_between_rank: np.ndarray,
rev_expert_per_rank: defaultdict[int, set[int]],
expert_from_rank: np.ndarray,
) -> tuple[np.ndarray, np.ndarray]:
"""
reselect and assign redundant experts to the ranks
with remaining redundant slots
"""
update_workload, num_per_existing_expert = self.recomputing_initial_weight(initial_weights, rank_assignments)
for rank_idx in undeployed_ranks:
for pos in redundant_expert_pos[rank_idx]:
sorted_expert_idx = np.argsort(update_workload, kind="stable")[::-1]
for expert_id in sorted_expert_idx:
send_rank = expert_from_rank[expert_id]
if expert_id in rank_assignments[rank_idx]:
continue
if np.isclose(update_workload[expert_id], -1):
raise ValueError(f"Expert ID {expert_id} is not in the node")
rank_assignments[rank_idx][pos] = expert_id
num_com_between_rank[send_rank][rank_idx] += 1
rev_expert_per_rank[rank_idx].add(expert_id)
num_cur_expert = num_per_existing_expert[expert_id]
update_workload[expert_id] *= num_cur_expert / (num_cur_expert + 1)
num_per_existing_expert[expert_id] += 1
break
rank_loads = np.zeros(len(rank_assignments), dtype=np.float32)
for rank_id, rank in enumerate(rank_assignments):
for index, expert_id in enumerate(rank):
rank_loads[rank_id] += update_workload[expert_id]
return update_workload, rank_loads
def non_redundant_expert_information(
self,
origin_deployment: np.ndarray,
updated_weights: np.ndarray,
redundant_expert_pos: list[list[int]],
) -> tuple[np.ndarray, np.ndarray]:
"""
Statistics on the status of logical experts on each rank
Parameters:
origin_deployment: [num_ranks, num_experts_per_rank]
the expert deployment status on each rank
updated_weights: [num_logical_experts] expert workload status after
reconfiguring redundant experts
redundant_expert_pos: the positions of redundant experts
on each rank
Returns:
rank_assignments: [num_ranks, num_experts_per_rank]
the deployment of logical experts on each rank
rank_loads: [num_ranks] The workload of
logical experts on each rank
"""
num_cur_deployment_ranks = origin_deployment.shape[0]
rank_assignments = np.full((num_cur_deployment_ranks, self.num_experts_per_rank), fill_value=-1, dtype=np.int64)
rank_loads = np.zeros(num_cur_deployment_ranks, dtype=np.float32)
for rank_id, rank in enumerate(origin_deployment):
for index, expert_id in enumerate(rank):
if index in redundant_expert_pos[rank_id]:
continue
rank_assignments[rank_id][index] = expert_id
rank_loads[rank_id] += updated_weights[expert_id]
return rank_assignments, rank_loads
def recomputing_initial_weight(
self, initial_weights: list, rank_assignments: np.ndarray
) -> tuple[np.ndarray, np.ndarray]:
"""
Calculate the load of the logic expert again based
on the current deployment
"""
num_per_existing_expert = np.zeros(self.num_original_experts, dtype=np.int64)
for rank in rank_assignments:
for expert_id in rank:
if expert_id != -1:
num_per_existing_expert[expert_id] += 1
update_workload = np.full(self.num_original_experts, fill_value=-1, dtype=np.float32)
for expert_id, weight in initial_weights:
num_cur_expert = num_per_existing_expert[expert_id]
assert num_cur_expert != 0
update_workload[expert_id] = weight / num_cur_expert
return update_workload, num_per_existing_expert
def distribute_redundant_experts(
self,
rank_assignments: np.ndarray,
rank_loads: np.ndarray,
redundant_expert_list: list[tuple[int, float]],
expert_from_rank: np.ndarray,
redundant_expert_pos: list[list[int]],
) -> tuple[np.ndarray, defaultdict[int, set[int]], list[int]]:
"""
Assign redundant experts to ranks
Parameters:
rank_assignments: [num_ranks, num_experts_per_rank]
the deployment of logical experts on each rank
rank_loads: [num_ranks] The workload of
logical experts on each rank
redundant_expert_list:[(expert_id, expert_load)]
the redundantly generated experts
expert_from_rank: [num_logical_experts] the rank where
each logical expert resides
redundant_expert_pos: the positions of redundant experts
on each rank
Returns:
num_com_between_rank:[num_ranks, num_ranks] the communication
status between ranks
rev_expert_per_rank:the experts assigned to each rank
after reconfiguring redundancy
undeployed_ranks: record the ranks that have not been
assigned redundant experts in redundancy positions
"""
num_ranks = len(rank_assignments)
rev_expert_per_rank = defaultdict(set)
num_com_between_rank = np.zeros((num_ranks, num_ranks), dtype=np.int64)
for expert_id, weight in redundant_expert_list:
candidate = -1
send_rank = expert_from_rank[expert_id]
for rank_id in range(num_ranks):
if len(redundant_expert_pos[rank_id]) == 0:
continue
if expert_id in rank_assignments[rank_id]:
continue
if num_com_between_rank[send_rank][rank_id] >= self.num_max_com:
continue
if candidate == -1 or rank_loads[rank_id] < rank_loads[candidate]:
candidate = rank_id
if candidate != -1:
pos = redundant_expert_pos[candidate].pop()
rank_assignments[candidate][pos] = expert_id
rank_loads[candidate] += weight
num_com_between_rank[send_rank][candidate] += 1
rev_expert_per_rank[candidate].add(expert_id)
undeployed_ranks = []
for rank_id in range(num_ranks):
if len(redundant_expert_pos[rank_id]) > 0:
undeployed_ranks.append(rank_id)
return num_com_between_rank, rev_expert_per_rank, undeployed_ranks
def redundancy_again(
self, cur_layer_workload: np.ndarray, cur_layer_deployment: np.ndarray
) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, defaultdict[int, set[int]]]:
"""
Calculate the status of a single node after redundant expert
configuration.
Parameters:
cur_layer_workload: [num_logical_experts] expert load statistics
cur_layer_deployment: [num_ranks, num_experts_per_rank]
the expert deployment status on each rank
Returns:
rank_assignments: [num_ranks, num_experts_per_rank]
the expert deployment status on each rank after
reconfiguring redundant experts
rank_loads: [num_ranks] the workload status of each rank after
reconfiguring redundant experts
updated_weights: [num_logical_experts] expert workload status after
reconfiguring redundant experts
num_com_between_rank:[num_ranks, num_ranks] the communication
status between ranks
rev_expert_per_rank:the experts assigned to each rank
after reconfiguring redundancy
"""
num_ranks = cur_layer_deployment.shape[0]
(redundant_expert_pos, expert_from_rank, existing_experts, num_redundant_experts) = (
self.statistics_expert_distribution(cur_layer_deployment)
)
initial_weights = []
for expert_id in existing_experts:
initial_weights.append((expert_id, cur_layer_workload[expert_id]))
redundant_expert_list, updated_weights = self.compute_redundant_assignments(
initial_weights, num_redundant_experts, num_ranks
)
rank_assignments, rank_loads = self.non_redundant_expert_information(
cur_layer_deployment, updated_weights, redundant_expert_pos
)
num_com_between_rank, rev_expert_per_rank, undeployed_ranks = self.distribute_redundant_experts(
rank_assignments, rank_loads, redundant_expert_list, expert_from_rank, redundant_expert_pos
)
if len(undeployed_ranks) > 0:
updated_weights, rank_loads = self.fill_in_undeployed_ranks(
initial_weights,
rank_assignments,
undeployed_ranks,
redundant_expert_pos,
num_com_between_rank,
rev_expert_per_rank,
expert_from_rank,
)
return (rank_assignments, rank_loads, updated_weights, num_com_between_rank, rev_expert_per_rank)
def redundant_expert_deployment(
self, cur_layer_workload: np.ndarray, cur_layer_deployment: np.ndarray
) -> tuple[
list[np.ndarray], list[np.ndarray], list[np.ndarray], list[np.ndarray], list[defaultdict[int, set[int]]]
]:
"""
Calculate the status of each node after reconfiguring redundant experts;
treat non-intra-node redundancy as a single node, store the result of
each node in a list, and return it.
"""
all_node_assignments = []
all_node_loads = []
updated_weights = []
num_com_between_rank = []
rev_expert_per_rank = []
if self.is_node_redundant:
num_ranks_per_node = self.num_die_per_host
for node_id in range(self.num_nodes):
cur_node_deployment = cur_layer_deployment[
node_id * num_ranks_per_node : (node_id + 1) * num_ranks_per_node
]
(
cur_node_rank_assignments,
cur_node_rank_loads,
cur_node_updated_weights,
cur_node_num_com_between_rank,
cur_node_rev_expert_per_rank,
) = self.redundancy_again(cur_layer_workload, cur_node_deployment)
all_node_assignments.append(cur_node_rank_assignments)
all_node_loads.append(cur_node_rank_loads)
updated_weights.append(cur_node_updated_weights)
num_com_between_rank.append(cur_node_num_com_between_rank)
rev_expert_per_rank.append(cur_node_rev_expert_per_rank)
else:
(
cur_rank_assignments,
cur_rank_loads,
cur_updated_weights,
cur_num_com_between_rank,
cur_rev_expert_per_rank,
) = self.redundancy_again(cur_layer_workload, cur_layer_deployment)
all_node_assignments.append(cur_rank_assignments)
all_node_loads.append(cur_rank_loads)
updated_weights.append(cur_updated_weights)
num_com_between_rank.append(cur_num_com_between_rank)
rev_expert_per_rank.append(cur_rev_expert_per_rank)
return (all_node_assignments, all_node_loads, updated_weights, num_com_between_rank, rev_expert_per_rank)
def swap_experts_between_ranks(
self,
max_rank_deployment_set: set[int],
swap_rank_deployment_set: set[int],
max_rank_rev_expert: set[int],
swap_rank_rev_expert: set[int],
workload: np.ndarray,
max_rank_load: float,
swap_rank_load: float,
) -> tuple[int, int, float]:
"""
Find the optimal experts for workload reduction
after exchange between two ranks
"""
max_rank_expert = -1
swap_rank_expert = -1
max_weight = max_rank_load
for cur_expert_id in max_rank_deployment_set:
if cur_expert_id in swap_rank_deployment_set or cur_expert_id in max_rank_rev_expert:
continue
cur_weight = float(workload[cur_expert_id])
for next_expert_id in swap_rank_deployment_set:
if next_expert_id in max_rank_deployment_set or next_expert_id in swap_rank_rev_expert:
continue
next_weight = float(workload[next_expert_id])
cur_load_after_swap = max_rank_load - cur_weight + next_weight
next_load_after_swap = swap_rank_load - next_weight + cur_weight
max_load_after_swap = max(cur_load_after_swap, next_load_after_swap)
if max_load_after_swap < max_weight:
max_weight = max_load_after_swap
max_rank_expert = cur_expert_id
swap_rank_expert = next_expert_id
return max_rank_expert, swap_rank_expert, max_weight
def expert_exchange_between_ranks(
self,
rank_assignments: np.ndarray,
rank_loads: np.ndarray,
num_com_between_rank: np.ndarray,
rev_expert_per_rank: defaultdict[int, set[int]],
updated_weights: np.ndarray,
) -> tuple[list[list[int]], float]:
"""
Perform inter-rank expert exchange within a single node
Parameters:
rank_assignments: [num_ranks, num_experts_per_rank] the expert
deployment status on each rank after reconfiguring
redundant experts
rank_loads: [num_ranks] the workload status of each rank after
reconfiguring redundant experts
num_com_between_rank:[num_ranks, num_ranks] the communication
status between ranks
rev_expert_per_rank:the experts assigned to each rank
after reconfiguring redundancy
updated_weights: [num_logical_experts] expert workload status after
reconfiguring redundant experts
Returns:
ranks_deployment_after_swap: [num_ranks, num_experts_per_rank]
the deployment status of experts on each rank after the exchange
max_rank_load: the workload of the hottest rank
"""
rank_deploy_sets = []
for rank_id in range(len(rank_assignments)):
rank_deploy_sets.append(set(rank_assignments[rank_id]))
max_swap_times = self.max_swap_times
max_rank_load = 0.0
exchange = True
while max_swap_times > 0:
max_swap_times -= 1
sorted_rank_idx = np.argsort(rank_loads, kind="stable")
max_load_rank_id = int(sorted_rank_idx[-1])
max_rank_load = float(rank_loads[max_load_rank_id])
if not exchange:
break
exchange = False
for swap_rank_id in sorted_rank_idx[:-1]:
if (
num_com_between_rank[swap_rank_id][max_load_rank_id] < self.num_max_com
and num_com_between_rank[max_load_rank_id][swap_rank_id] < self.num_max_com
):
swap_rank_load = rank_loads[swap_rank_id]
max_rank_expert, swap_rank_expert, max_weight = self.swap_experts_between_ranks(
rank_deploy_sets[max_load_rank_id],
rank_deploy_sets[swap_rank_id],
rev_expert_per_rank[max_load_rank_id],
rev_expert_per_rank[swap_rank_id],
updated_weights,
max_rank_load,
swap_rank_load,
)
if max_rank_load - max_weight < self.swap_threshold or max_rank_expert == -1:
continue
rank_deploy_sets[max_load_rank_id].remove(max_rank_expert)
rank_deploy_sets[swap_rank_id].remove(swap_rank_expert)
rank_deploy_sets[max_load_rank_id].add(swap_rank_expert)
rank_deploy_sets[swap_rank_id].add(max_rank_expert)
rank_loads[max_load_rank_id] += updated_weights[swap_rank_expert] - updated_weights[max_rank_expert]
rank_loads[swap_rank_id] += updated_weights[max_rank_expert] - updated_weights[swap_rank_expert]
rev_expert_per_rank[max_load_rank_id].add(swap_rank_expert)
rev_expert_per_rank[swap_rank_id].add(max_rank_expert)
num_com_between_rank[swap_rank_id][max_load_rank_id] += 1
num_com_between_rank[max_load_rank_id][swap_rank_id] += 1
exchange = True
break
ranks_deployment_after_swap = [list(s) for s in rank_deploy_sets]
return ranks_deployment_after_swap, max_rank_load
def exchange_experts(
self,
all_node_assignments: list[np.ndarray],
all_node_loads: list[np.ndarray],
num_com_between_rank: list[np.ndarray],
rev_expert_per_rank: list[defaultdict[int, set[int]]],
updated_weights: list[np.ndarray],
) -> tuple[np.ndarray, float]:
"""
For each node after redundancy, perform inter-rank expert
exchange within the node to reduce the workload of the hottest rank
"""
max_workload = 0.0
after_swap_ranks_deployment = []
for idx in range(len(all_node_assignments)):
cur_node_deployment, cur_node_max_workload = self.expert_exchange_between_ranks(
all_node_assignments[idx],
all_node_loads[idx],
num_com_between_rank[idx],
rev_expert_per_rank[idx],
updated_weights[idx],
)
after_swap_ranks_deployment += cur_node_deployment
if cur_node_max_workload > max_workload:
max_workload = cur_node_max_workload
new_deployment = np.array(after_swap_ranks_deployment)
return new_deployment, max_workload
def rebalance_experts(
self, current_expert_table: torch.Tensor, expert_workload: torch.Tensor, is_node_redundant: bool = False
) -> tuple[int, np.ndarray, list[list[list[int]]]]:
"""
Rebalance experts based on workload and deployment strategy.
Parameters:
current_expert_table: [num_layer, num_ranks, num_experts_per_rank],
tensor of current expert assignment.
expert_workload: [num_layer, num_ranks, num_experts_per_rank],
tensor of workload for each expert.
is_node_redundant: Whether to enable intra-node redundancy.
Returns:
change: Scalar flag for whether deployment changed.
per_layer_priority: List of priority for each layer.
new_deployment: [num_layer, num_ranks, num_experts_per_rank],
list of adjusted expert deployment result.
"""
info = DynamicTable()
info.workload_table = expert_workload.numpy()
info.placement_table = current_expert_table.numpy()
assert info.workload_table is not None and info.placement_table is not None
self.is_node_redundant = is_node_redundant
self.num_layers, self.num_ranks, self.num_experts_per_rank = info.placement_table.shape
self.num_nodes = self.num_ranks // self.num_die_per_host
expert_ids, counts = np.unique(info.placement_table[0], return_counts=True)
self.num_original_experts = len(expert_ids)
layer_workloads = self.get_original_workload(
info.placement_table, info.workload_table, self.num_original_experts
)
per_layer_total_load = layer_workloads[0].sum()
ave_workload = per_layer_total_load / self.num_ranks
self.swap_threshold = ave_workload * self.increment
layer_initial_imbalance = self.calculate_imbalance(info.placement_table, layer_workloads)
new_deployment = info.placement_table.copy()
max_heat_per_layer_before = self.calculate_max_heat_per_layer(info.workload_table)
npu_heat_all_origin = sum(max_heat_per_layer_before)
max_heat_per_layer_after = []
for layer in range(self.num_layers):
cur_layer_deployment = info.placement_table[layer]
cur_layer_workload = layer_workloads[layer]
if layer_initial_imbalance[layer] < self.imbalance_threshold:
max_heat_per_layer_after.append(max_heat_per_layer_before[layer])
continue
(all_node_assignments, all_node_loads, updated_weights, num_com_between_rank, rev_experts_per_rank) = (
self.redundant_expert_deployment(cur_layer_workload, cur_layer_deployment)
)
(new_layer_deployment, new_max_workload) = self.exchange_experts(
all_node_assignments, all_node_loads, num_com_between_rank, rev_experts_per_rank, updated_weights
)
after_swap_imbalance = new_max_workload / ave_workload
if after_swap_imbalance < layer_initial_imbalance[layer]:
new_deployment[layer] = new_layer_deployment
max_heat_per_layer_after.append(new_max_workload)
self.constraint_expert_local_exchange(info.placement_table, new_deployment)
layer_changed_ratio = []
for layer_idx in range(self.num_layers):
if max_heat_per_layer_before[layer_idx] > 0:
layer_changed_ratio.append(max_heat_per_layer_after[layer_idx] / max_heat_per_layer_before[layer_idx])
else:
layer_changed_ratio.append(1.0)
per_layer_priority = np.argsort(layer_changed_ratio)
npu_heat_all_after = sum(max_heat_per_layer_after)
change = 0
if npu_heat_all_after < 0.95 * npu_heat_all_origin:
change = 1
return change, per_layer_priority, new_deployment.tolist()

187
vllm_ascend/eplb/eplb_updator.py
View File

@@ -19,181 +19,168 @@ import numpy
import torch
import torch.distributed as dist
import vllm.envs as envs
from vllm.distributed.parallel_state import get_pp_group
from vllm.logger import logger
from vllm.v1.utils import record_function_or_nullcontext
from vllm_ascend.distributed.parallel_state import get_dynamic_eplb_group
from vllm_ascend.eplb.adaptor.vllm_adaptor import VllmEplbAdaptor
from vllm_ascend.eplb.core.eplb_device_transfer_loader import D2DExpertWeightLoader
from vllm_ascend.eplb.core.eplb_worker import EplbProcess
class EplbUpdator:
def __init__(self, ascend_config, loader, eplb_process: EplbProcess,
process):
self.ascend_config = ascend_config
self.init_eplb(self.ascend_config.expert_map_path, process)
def __init__(self, eplb_config, loader: D2DExpertWeightLoader, eplb_process: EplbProcess, process):
self.eplb_config = eplb_config
self.multi_stage = eplb_config.eplb_policy_type == 3
self.init_eplb(self.eplb_config.expert_map_path, process)
self.eplb_loader = loader
self.eplb_process = eplb_process
self.shared_dict = self.eplb_process.shared_dict
self.comm_group = get_dynamic_eplb_group()
def set_adaptor(self, adaptor):
def set_adaptor(self, adaptor: VllmEplbAdaptor):
self.pp_rank = get_pp_group().rank_in_group
self.adaptor = adaptor
self.num_moe_layers = self.adaptor.num_moe_layers
self.global_expert_num = self.adaptor.global_expert_num
local_load = self.adaptor.get_rank_expert_workload()
self.world_size = dist.get_world_size()
self.device = local_load.device
self.eplb_loader.num_layers = self.adaptor.num_dense_layers + self.adaptor.num_moe_layers
def init_eplb(self, expert_map_path, process):
self.rank_id = dist.get_rank()
self.num_expert_load_gather = 10
self.periodic_load_gather = True
self.num_iterations_eplb_update: torch.int64 = self.ascend_config.num_iterations_eplb_update
self.expert_heat_collection_interval: torch.int64 = self.eplb_config.expert_heat_collection_interval
self.expert_map_path = expert_map_path
self.expert_map_record_path = self.ascend_config.expert_map_record_path
self.expert_map_record_path = self.eplb_config.expert_map_record_path
try:
if not envs.VLLM_ALLOW_EXPERT_LOAD_COLLECTING:
self.num_expert_load_gather = self.num_iterations_eplb_update
self.num_expert_load_gather = self.expert_heat_collection_interval
self.periodic_load_gather = False
except Exception:
self.num_expert_load_gather = self.num_iterations_eplb_update
logger.debug("[eplb/updator] VLLM_ALLOW_EXPERT_LOAD_COLLECTING unavailable in current vllm version.")
self.num_expert_load_gather = self.expert_heat_collection_interval
self.periodic_load_gather = False
self.expert_map_initialized = False
self.gate_eplb = self.ascend_config.gate_eplb
self.reqs = []
self.update_info_all = []
self.cur_iterations: torch.int64 = 0
self.num_wait_worker_iterations: torch.int64 = self.ascend_config.num_wait_worker_iterations
self.algorithm_execution_interval: torch.int64 = self.eplb_config.algorithm_execution_interval
self.process = process
logger.info(
f"[ModelRunner] Launched EPLB process (pid={self.process.pid})")
logger.info("[eplb/updator] Launched EPLB subprocess, pid=%s", self.process.pid)
def update_iteration(self):
self.cur_iterations += 1
if self.cur_iterations == (self.num_iterations_eplb_update + \
self.num_wait_worker_iterations + self.num_moe_layers):
if self.cur_iterations == (
self.expert_heat_collection_interval + self.algorithm_execution_interval + self.num_moe_layers
):
logger.debug("[eplb/updator] Full EPLB cycle completed, clearing moe loads and resetting iteration counter")
if self.expert_map_record_path is not None:
self.adaptor._export_tensor_to_file(
self.shared_dict["expert_maps"],
self.expert_map_record_path)
self.adaptor._export_tensor_to_file(self.shared_dict["expert_maps"], self.expert_map_record_path)
self.adaptor.model.clear_all_moe_loads()
if not self.gate_eplb:
self.cur_iterations = 0
self.adaptor.clear_all_moe_loads()
self.cur_iterations = 0
def get_update_info_flag(self):
return self.cur_iterations == (self.num_iterations_eplb_update +
self.num_wait_worker_iterations - 1)
return self.cur_iterations == (self.expert_heat_collection_interval + self.algorithm_execution_interval - 1)
def wakeup_eplb_worker_flag(self):
return self.cur_iterations == (self.num_iterations_eplb_update - 1)
return self.cur_iterations == (self.expert_heat_collection_interval - 1)
def update_expert_weight_flag(self):
weight_update_counter = self.cur_iterations - (
self.num_iterations_eplb_update + self.num_wait_worker_iterations)
return (weight_update_counter >= 0
and weight_update_counter < self.num_moe_layers)
def get_init_expert_map(self):
try:
if not self.expert_map_initialized:
self.shared_dict[
"expert_maps"] = self.adaptor.get_init_expert_map_from_file(
self.num_moe_layers, self.expert_map_path)
self.expert_map_initialized = True
except Exception as e:
logger.warning(f"[ModelRunner] Failed to wake EPLB process: {e}",
exc_info=True)
self.expert_heat_collection_interval + self.algorithm_execution_interval
)
return weight_update_counter >= 0 and weight_update_counter < self.num_moe_layers
def wakeup_eplb_worker(self):
self.eplb_process.planner_q.put(1)
def forward_before(self):
if self.update_expert_weight_flag():
(expert_send_info, expert_recv_info, updated_expert_map,
log2phy_map, layer_id) = self.update_info_all.pop(0)
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))
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)
# set asynchronous stream for d2d expert weight update
self.reqs = []
self.eplb_loader.asyn_expert_weight_transfer(self.reqs)
def take_update_info_from_eplb_process(self):
# Batch after eplb process being triggered, get update info provided by eplb process
if self.get_update_info_flag():
self.update_info_all = self.eplb_process.block_update_q.get()
def forward_end(self):
if self.wakeup_eplb_worker_flag():
self.compute_and_set_moe_load(is_clear=True)
self.wakeup_eplb_worker()
if self.update_expert_weight_flag():
with record_function_or_nullcontext("EPLB generate p2p task"):
(expert_send_info, expert_recv_info, updated_expert_map, log2phy_map, layer_id) = (
self.update_info_all.pop(0)
)
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))
self.eplb_loader.generate_expert_d2d_transfer_task(
expert_send_info,
expert_recv_info,
updated_expert_map_this_rank,
layer_id,
)
# set asynchronous stream for d2d expert weight update
self.reqs = []
self.eplb_loader.asyn_expert_weight_transfer(self.reqs)
def forward_end(self, eplb_heat_collection_status: bool = True):
if self.wakeup_eplb_worker_flag():
with record_function_or_nullcontext("EPLB gather moe load"):
self.compute_and_set_moe_load()
self.wakeup_eplb_worker()
if self.update_expert_weight_flag() and self.expert_map_record_path is None:
self.eplb_loader.update_expert_map_and_weight(self.reqs)
self.update_iteration()
# One circle of eplb update includes expert_heat_collection_interval + algorithm_execution_interval
# + num_moe_layers (for weight update). In expert_heat_collection stage, we only update the counter
# when eplb_heat_collection_status is True. In later stages, the counter is always updated.
# TODO(Angazenn): Decouple algorithm execution && weight update with heat collection iterations.
if self.cur_iterations >= self.expert_heat_collection_interval - 1 or eplb_heat_collection_status:
self.update_iteration()
def compute_and_set_moe_load(self, is_clear=False):
local_load = self.adaptor.get_rank_expert_workload()
def compute_and_set_moe_load(self):
local_load = self.adaptor.get_rank_expert_workload().unsqueeze(1)
moe_load = self.comm_group.all_gather(local_load, dim=1).cpu()
self._gather_buffer = None
if dist.is_initialized():
self.world_size = dist.get_world_size()
self.device = local_load.device
if self._gather_buffer is None:
shape = (self.world_size, *local_load.shape)
self._gather_buffer = torch.empty(shape,
dtype=local_load.dtype,
device=self.device)
if self.multi_stage:
moe_load = moe_load.permute(2, 0, 1, 3)
dist.all_gather_into_tensor(self._gather_buffer, local_load)
self.shared_dict["moe_load"] = moe_load
logger.debug("[eplb/updator] Updated shared_dict['moe_load'] shape=%s", moe_load.shape)
moe_load = self._gather_buffer.permute(1, 0, 2)
self.shared_dict["moe_load"] = moe_load.cpu()
logger.debug(
f"[ModelRunner] Updated shared_dict['moe_load'] shape={moe_load.shape}"
)
else:
moe_load = local_load.unsqueeze(1)
self.shared_dict["moe_load"] = moe_load.cpu()
logger.debug(
f"[ModelRunner] Updated shared_dict['moe_load'] shape={moe_load.shape}"
)
return moe_load
def warm_up_eplb(self):
self.get_init_expert_map()
logger.info("[eplb/updator] Starting EPLB warm-up, rank=%s, world_size=%s", self.rank_id, self.world_size)
self.shared_dict["expert_maps"] = self.adaptor.get_global_expert_map()
self.compute_and_set_moe_load()
src_tensor = torch.empty((1, ), device=self.device)
self_rank = dist.get_rank()
src_tensor = torch.empty((1,), device=self.device)
comm_op_list = []
reqs = []
for dst_rank in range(self.world_size):
if dst_rank == self_rank:
for dst_rank in range(self.comm_group.world_size):
if dst_rank == self.comm_group.rank_in_group:
continue
comm_op_list.append(dist.P2POp(dist.isend, src_tensor, dst_rank))
global_dst = self.comm_group.ranks[dst_rank]
comm_op_list.append(dist.P2POp(dist.isend, src_tensor, global_dst, group=self.comm_group.device_group))
for src_rank in range(self.world_size):
if src_rank == self_rank:
for src_rank in range(self.comm_group.world_size):
if src_rank == self.comm_group.rank_in_group:
continue
comm_op_list.append(dist.P2POp(dist.irecv, src_tensor, src_rank))
global_src = self.comm_group.ranks[src_rank]
comm_op_list.append(dist.P2POp(dist.irecv, src_tensor, global_src, group=self.comm_group.device_group))
if comm_op_list:
reqs = dist.batch_isend_irecv(comm_op_list)
for req in reqs:
req.wait()
logger.info("[eplb/updator] EPLB warm-up completed")
def shutdown(self):
"""
@@ -202,4 +189,4 @@ class EplbUpdator:
if self.process.is_alive():
self.process.terminate()
self.process.join()
logger.info("[ModelRunner] EPLB process terminated")
logger.info("[eplb/updator] EPLB subprocess terminated")