forked from EngineX-Cambricon/enginex-mlu370-vllm
144 lines
9.1 KiB
Python
144 lines
9.1 KiB
Python
import torch
|
|
import torch_mlu
|
|
import torch_mlu_ops as tmo
|
|
from common import benchmark_forward, save_to_csv
|
|
import argparse
|
|
from tabulate import tabulate
|
|
import os
|
|
|
|
e2e_time_param_dict_list = [
|
|
{"batch": 1, "seq_len": 1, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
|
|
"gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
|
|
"topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},
|
|
|
|
{"batch": 16, "seq_len": 1, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
|
|
"gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
|
|
"topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},
|
|
|
|
{"batch": 72, "seq_len": 1, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
|
|
"gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
|
|
"topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},
|
|
|
|
{"batch": 128, "seq_len": 1, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
|
|
"gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
|
|
"topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},
|
|
|
|
{"batch": 490, "seq_len": 1, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
|
|
"gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
|
|
"topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},
|
|
|
|
{"batch": 525, "seq_len": 1, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
|
|
"gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
|
|
"topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},
|
|
|
|
{"batch": 1, "seq_len": 1024, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
|
|
"gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
|
|
"topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},
|
|
|
|
{"batch": 1, "seq_len": 2048, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
|
|
"gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
|
|
"topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},
|
|
|
|
{"batch": 1, "seq_len": 4096, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
|
|
"gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
|
|
"topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},
|
|
|
|
{"batch": 1, "seq_len": 8192, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
|
|
"gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
|
|
"topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},
|
|
|
|
{"batch": 1, "seq_len": 32768, "hidden_size": 8192, "inner_size": 1024, "num_expert": 32, "start_expert_id": 0, "expert_size": 32,
|
|
"gated_ffn": False, "has_residual": False, "smooth_quant": True, "act_mode": "gelu",
|
|
"topk": 5, "renormalize": False, "dtype": [torch.bfloat16]},
|
|
]
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--repeat_times', type=int, default=10, help='repeat times for testing')
|
|
parser.add_argument('--csv', action='store_true', help='write the report data to csv')
|
|
parser.add_argument('-o', type=str, help='specify the output folder name under --csv mode')
|
|
|
|
args = parser.parse_args()
|
|
device = 'mlu'
|
|
|
|
titles = ["batch", "seq_len", "hidden_size", "inner_size", "gated_ffn", "num_expert", "topk", "act_mode", "quant_weight", "dtype", "hardware_time(us)", "e2e_latency(us)"]
|
|
contents = []
|
|
for params_dict in e2e_time_param_dict_list:
|
|
batch = params_dict["batch"]
|
|
seq_len = params_dict["seq_len"]
|
|
hidden_size = params_dict["hidden_size"]
|
|
inner_size = params_dict["inner_size"]
|
|
gated_ffn = params_dict["gated_ffn"]
|
|
act_mode = params_dict["act_mode"]
|
|
num_expert = params_dict["num_expert"]
|
|
start_expert_id = params_dict["start_expert_id"]
|
|
expert_size = params_dict["expert_size"]
|
|
topk = params_dict["topk"]
|
|
has_residual = params_dict["has_residual"]
|
|
smooth_quant = params_dict["smooth_quant"]
|
|
renormalize = params_dict["renormalize"]
|
|
input_dtype_list = params_dict["dtype"]
|
|
# print(f"batch:{batch}, seq_len:{seq_len}, hidden_size:{hidden_size}, inner_size:{inner_size}, "
|
|
# f"gated_ffn:{gated_ffn}, act_mode:{act_mode}, num_expert:{num_expert}, topk:{topk}")
|
|
for dtype in input_dtype_list:
|
|
if dtype == torch.bfloat16 and not torch_mlu.mlu.is_bf16_supported():
|
|
dtype = torch.half
|
|
hidden_states = torch.randn(batch, seq_len, hidden_size, device=device, dtype=dtype)
|
|
router_logit = torch.randn(batch, seq_len, num_expert, device=device, dtype=torch.float32)
|
|
|
|
if False: # print token_count
|
|
softmax = torch.softmax(router_logit.view(-1, router_logit.size(-1)), dim=1)
|
|
topk_logit, expert_id = torch.topk(softmax, k=topk, dim=1)
|
|
if renormalize:
|
|
topk_logit = topk_logit / topk_logit.sum(-1).unsqueeze(1)
|
|
sorted_expert_id, indices = expert_id.int().flatten().sort()
|
|
token_cout = torch.bincount(sorted_expert_id, minlength=num_expert).int()
|
|
print(token_cout)
|
|
|
|
residual = None
|
|
if has_residual:
|
|
residual = torch.randn(batch, seq_len, hidden_size, device=device, dtype=dtype)
|
|
weight1 = torch.randn(num_expert, inner_size*(1+gated_ffn), hidden_size, device=device, dtype=dtype)
|
|
bias1 = None # torch.randn(expert_num, inner_size*(1+gated), device=device, dtype=data_type)
|
|
weight2 = torch.randn(num_expert, hidden_size, inner_size, device=device, dtype=dtype)
|
|
bias2 = None # torch.randn(expert_num, hidden_size, device=device, dtype=data_type)
|
|
input_smooth, act_smooth, w1_scale, w2_scale = None, None, None, None
|
|
if smooth_quant:
|
|
input_smooth = torch.randn(expert_size, hidden_size, device=device, dtype=torch.float32).abs() + 0.1
|
|
act_smooth = torch.randn(expert_size, inner_size, device=device, dtype=torch.float32).abs() + 0.1
|
|
weight1 = torch.randint(-128, 127, (num_expert, inner_size*(1+gated_ffn), hidden_size)).to(torch.int8).mlu()
|
|
weight2 = torch.randint(-128, 127, (num_expert, hidden_size, inner_size)).to(torch.int8).mlu()
|
|
w1_scale = torch.randn(expert_size, (1+gated_ffn)*inner_size).to(device).to(torch.float32)
|
|
w2_scale = torch.randn(expert_size, hidden_size).to(device).to(torch.float32)
|
|
|
|
hardware_time, e2e_time = benchmark_forward(tmo.fused_moe,
|
|
hidden_states,
|
|
router_logit,
|
|
weight1[start_expert_id:start_expert_id+expert_size],
|
|
weight2[start_expert_id:start_expert_id+expert_size],
|
|
bias1,
|
|
bias2,
|
|
residual,
|
|
input_smooth,
|
|
act_smooth,
|
|
w1_scale,
|
|
w2_scale,
|
|
topk,
|
|
renormalize,
|
|
gated_ffn,
|
|
act_mode,
|
|
start_expert_id,
|
|
repeats=args.repeat_times)
|
|
content = [f"{batch}", f"{seq_len}", f"{hidden_size}", f"{inner_size}", f"{gated_ffn}", f"{num_expert}", f"{topk}", f"{act_mode}", f"{smooth_quant}", f"{dtype}", f"{hardware_time}", f"{e2e_time}"]
|
|
contents.append(content)
|
|
table = [titles] + contents
|
|
print(tabulate(table, headers="firstrow", tablefmt="grid"))
|
|
|
|
if args.csv:
|
|
current_file_path = __file__
|
|
_, file_name = os.path.split(current_file_path)
|
|
save_to_csv(table, args.o, file_name)
|
|
|
|
if __name__=="__main__":
|
|
main()
|