187 lines
6.8 KiB
Python
187 lines
6.8 KiB
Python
# coding=utf-8
|
||
# Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
|
||
import json
|
||
import logging
|
||
import os
|
||
|
||
import matplotlib.pyplot as plt # type: ignore
|
||
import numpy as np
|
||
import torch
|
||
|
||
os.environ["VLLM_USE_MODELSCOPE"] = "True"
|
||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||
|
||
logger = logging.getLogger("msit_logger")
|
||
|
||
|
||
def save_matrix_to_json(output_path, file_name, deployment):
|
||
num_layers = deployment.shape[0]
|
||
num_cards = deployment.shape[1]
|
||
|
||
data = {"moe_layer_count": num_layers}
|
||
layer_list = []
|
||
for i in range(num_layers):
|
||
layer = {"layer_id": i, "device_count": num_cards}
|
||
device_list = []
|
||
for j in range(num_cards):
|
||
device = {
|
||
"device_id": j,
|
||
"device_expert": deployment[i, j].tolist()
|
||
}
|
||
device_list.append(device)
|
||
layer["device_list"] = device_list
|
||
layer_list.append(layer)
|
||
data["layer_list"] = layer_list
|
||
|
||
file_name = f"{output_path}{file_name}.json"
|
||
|
||
# Save as JSON file
|
||
try:
|
||
with open(file_name, 'w') as f:
|
||
json.dump(data, f, indent=4)
|
||
except Exception as e:
|
||
print(f"write {file_name} failed: {e}")
|
||
|
||
|
||
def calculate_average(lst):
|
||
"""calculate the average of a list"""
|
||
if not lst:
|
||
raise ValueError("list is empty")
|
||
|
||
total = 0.0
|
||
count = 0
|
||
|
||
for element in lst:
|
||
# Check if element is numeric
|
||
if isinstance(element, (int, float, np.int64, np.float64)):
|
||
total += float(element)
|
||
count += 1
|
||
else:
|
||
# Non-numeric elements will be ignored with a warning
|
||
print(f"warning: element {element} is not a number, ignored")
|
||
|
||
if count == 0:
|
||
raise ValueError("list does not contain any number")
|
||
|
||
return total / count
|
||
|
||
|
||
def layer_imblance_polt(y_list, label_names, device_num, output_path,
|
||
file_name):
|
||
|
||
plt.rcParams['font.sans-serif'] = ['Arial']
|
||
plt.rcParams['axes.unicode_minus'] = False
|
||
x = [i for i in range(58)]
|
||
for index, y in enumerate(y_list):
|
||
plt.plot(x,
|
||
y,
|
||
label=rf'{label_names[index]}，avg={calculate_average(y)}')
|
||
|
||
plt.legend()
|
||
plt.title(rf'Load Distribution (num_gpus={device_num})')
|
||
plt.xlabel('layer')
|
||
plt.ylabel('Device Load')
|
||
|
||
# Show grid lines
|
||
plt.grid(True)
|
||
|
||
plt.savefig(os.path.join(output_path, file_name), dpi=300)
|
||
|
||
# Clear current plot
|
||
plt.close()
|
||
|
||
|
||
def deepseek_deploy(workload, num_redundancy_expert, num_groups, num_nodes,
|
||
num_gpus, num_original_expert):
|
||
from eplb_deepseek import rebalance_experts
|
||
num_replicas = num_original_expert + num_redundancy_expert
|
||
hy2log, log2phy, logcnt = rebalance_experts(workload, num_replicas,
|
||
num_groups, num_nodes,
|
||
num_gpus)
|
||
|
||
# Convert to global_deployment
|
||
workload = workload.cpu().numpy()
|
||
global_deployment = []
|
||
layer_num = log2phy.shape[0]
|
||
num_physical_experts_local = (num_original_expert +
|
||
num_redundancy_expert) // num_gpus
|
||
for layer_idx in range(layer_num):
|
||
layer_deployment = []
|
||
for gpu_idx in range(num_gpus):
|
||
local_deployment = hy2log[layer_idx][gpu_idx *
|
||
num_physical_experts_local:
|
||
(gpu_idx + 1) *
|
||
num_physical_experts_local]
|
||
local_deployment = local_deployment.flatten()
|
||
layer_deployment.append(local_deployment.tolist())
|
||
global_deployment.append(layer_deployment)
|
||
|
||
# Remap expert distribution according to log2phy
|
||
original_weights = []
|
||
max_weights = []
|
||
average_weights = []
|
||
y_list = []
|
||
for layer_idx in range(layer_num):
|
||
new_value = workload[layer_idx].reshape(num_gpus, -1)
|
||
row_sum = np.sum(new_value, axis=1)
|
||
original_weights.append(row_sum.max())
|
||
average_weights.append((np.sum(workload[layer_idx]) / num_gpus))
|
||
|
||
opt_workload = np.zeros((num_original_expert + num_redundancy_expert),
|
||
dtype=np.float64)
|
||
for expert_idx in range(num_original_expert):
|
||
physical_expert_idxs = log2phy[layer_idx][expert_idx]
|
||
physical_expert_idxs = physical_expert_idxs.flatten()
|
||
physical_expert_idxs = physical_expert_idxs[
|
||
physical_expert_idxs != -1]
|
||
for physical_expert_idx in physical_expert_idxs:
|
||
opt_workload[physical_expert_idx] += workload[layer_idx][
|
||
expert_idx] / len(physical_expert_idxs)
|
||
opt_workload = opt_workload.reshape(num_gpus, -1)
|
||
row_sum = np.sum(opt_workload, axis=1)
|
||
max_weights.append(row_sum.max())
|
||
|
||
y_list = [original_weights, max_weights, average_weights]
|
||
return global_deployment, y_list
|
||
|
||
|
||
if __name__ == '__main__':
|
||
import argparse
|
||
parser = argparse.ArgumentParser()
|
||
parser.add_argument("--exp_name", type=str, default="gsm8k_temp0.0")
|
||
parser.add_argument("--num_original_expert", type=int, default=256)
|
||
parser.add_argument("--input_path", type=str, default="")
|
||
parser.add_argument("--output_path", type=str, default="")
|
||
parser.add_argument("--num_redundancy_expert", type=int, default=0)
|
||
parser.add_argument("--num_devices", type=int, default=32)
|
||
parser.add_argument("--num_groups", type=int, default=8)
|
||
parser.add_argument("--num_nodes", type=int, default=4)
|
||
args = parser.parse_args()
|
||
exp_name = args.exp_name
|
||
input_path = args.input_path
|
||
output_path = args.output_path
|
||
os.makedirs(output_path, exist_ok=True)
|
||
num_redundancy_expert = args.num_redundancy_expert
|
||
num_devices = args.num_devices
|
||
num_original_expert = args.num_original_expert
|
||
num_groups = args.num_groups
|
||
num_nodes = args.num_nodes
|
||
|
||
# NOTE: assume input workload format: [layer_num, num_experts]
|
||
workload = torch.load(input_path, map_location=torch.device('cpu'))
|
||
global_deployment, y_list = deepseek_deploy(workload,
|
||
num_redundancy_expert,
|
||
num_groups, num_nodes,
|
||
num_devices,
|
||
num_original_expert)
|
||
|
||
file_name = f"{exp_name}_{num_devices}_{num_redundancy_expert}"
|
||
save_matrix_to_json(output_path, file_name, np.array(global_deployment))
|
||
label_names = [
|
||
'default deployment max load', 'balanced load max load',
|
||
'balanced load avg load'
|
||
]
|
||
new_file_name = f"{exp_name}_{num_devices}_{num_redundancy_expert}.png"
|
||
layer_imblance_polt(y_list, label_names, num_devices, output_path,
|
||
new_file_name)
|