@@ -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]
|
||||
|
||||
Reference in New Issue
Block a user