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()