import os import sys import time import pathlib import argparse from tqdm import tqdm import numpy as np import prettytable as pt import csv import importlib import json import random FILE_DIR = pathlib.Path(__file__).parent.absolute() # ${prj_root}/byte_infer_perf INFER_ROOT_DIR = FILE_DIR.parent CUR_DIR = pathlib.Path.cwd().absolute() os.chdir(INFER_ROOT_DIR) sys.path.insert(0, str(INFER_ROOT_DIR)) from llm_perf.utils.logger import logger, setup_logger def parse_args(): default_model_config = FILE_DIR.joinpath( "model_zoo", "mixtral-torch-bf16-8x22b.json" ).absolute() parser = argparse.ArgumentParser() parser.add_argument("--hardware_type", type=str, default="GPU") # model_info parser.add_argument("--model_config", type=str, default=default_model_config) # deploy_config parser.add_argument("--tp_size", type=int, default=8) parser.add_argument("--max_batch_size", type=int, default=8) # realization_version, mapping to: # 1. engine_type: tp_engine, pp_engine, ... # 2. model realization info, parser.add_argument("--model_version", type=str, default="default") parser.add_argument("--infer_dtype", type=str, default="bfloat16") # perf config, using default config parser.add_argument("--perf_config", type=str, default=None) # feature control parser.add_argument("--random_seed", type=int, default=1) parser.add_argument("--log_model", action="store_true") parser.add_argument("--workspace", type=str, default=None) parser.add_argument("--log_level", type=str, default="info") return parser.parse_args() def perf_engine(xpu_config, workspace): hardware_type = xpu_config["hardware_type"] setup = importlib.import_module( ".setup", package=f"llm_perf.backends.{hardware_type}" ) logger.info(f"import setup: {setup}") xpu_config["sku_name"] = setup.get_device_name() with open(workspace.joinpath("config.json"), "w") as f: json.dump(xpu_config, f, indent=4) random_seed = xpu_config["feature_config"]["random_seed"] log_model = xpu_config["feature_config"]["log_model"] max_batch_size = xpu_config["deploy_config"]["max_batch_size"] # create mp_engine (currently only support tp) # 1. tp_size # 2. max_batch_size engine = setup.get_engine(xpu_config) def update_template(mode, batch_size, seq_len): if mode == "context": input_template = { "is_context": True, "input_ids": [[random.randint(1000, 5000) for _ in range(seq_len)] for _ in range(batch_size)], "position_ids": [[i for i in range(seq_len)] for _ in range(batch_size)], "attention_mask": [[1 for _ in range(seq_len)] for _ in range(batch_size)], "valid_slot_ids": [i for i in range(batch_size)], "all_q_len": [seq_len for _ in range(batch_size)], "all_kv_len": [seq_len for _ in range(batch_size)], "get_input_logits": False, "workspace": workspace, "log": log_model, "override_hidden_states": True, "random_seed": random_seed } elif mode == "decode": input_template = { "is_context": False, "input_ids": [[random.randint(1000, 5000)] for _ in range(batch_size)], "position_ids": [[seq_len] for _ in range(batch_size)], "attention_mask": [[1] for _ in range(batch_size)], "valid_slot_ids": [i for i in range(batch_size)], "all_q_len": [1 for _ in range(batch_size)], "all_kv_len": [seq_len for _ in range(batch_size)], "get_input_logits": False, "workspace": workspace, "log": log_model, "override_hidden_states": True, "random_seed": random_seed } return input_template # warmup num_warm_iter = 10 input_template = update_template("context", 1, 1024) start_time = time.perf_counter_ns() for _ in range(num_warm_iter): engine.mp_forward(input_template) duration_s = round((time.perf_counter_ns() - start_time) / 1e9, 3) logger.info(f"warmup cost: {duration_s}s") def results_to_csv(file_path, results): batch_size_set = sorted(results.keys()) seq_len_set = set() for batch_size in batch_size_set: for seq_len in results[batch_size].keys(): seq_len_set.add(seq_len) seq_len_set = sorted(list(seq_len_set)) with open(file_path, "w") as f: csv_writer = csv.writer(f) header = ["batch_size\n----------\nseq_len"] for batch_size in batch_size_set: header.append(f"{batch_size}") csv_writer.writerow(header) for seq_len in seq_len_set: row = [f"{seq_len}"] for batch_size in batch_size_set: row.append(f"{results[batch_size][seq_len]}") csv_writer.writerow(row) log_results = [] if xpu_config["perf_config"]["perf_context"]: batch_size_list = [1] seq_len_list = xpu_config["perf_config"]["seq_len_list"] context_results = {} for batch_size in batch_size_list: context_results[batch_size] = {} for seq_len in seq_len_list: input_template = update_template("context", 1, seq_len) start_iters = 2 test_iter = 0 duration_ms = 0. while duration_ms < 5000. and test_iter < 100: result = engine.mp_forward(input_template) if start_iters > 0: start_iters -= 1 continue test_iter += 1 duration_ms += result["duration_ms"] avg_duration_ms = round(duration_ms / test_iter, 3) context_results[batch_size][seq_len] = avg_duration_ms log_results.append(f"context, batch_size={batch_size}, seq_len={seq_len}, avg_duration_ms={avg_duration_ms}") if log_model: log_path = workspace.joinpath("rank_0", "run.log") if log_path.exists(): lines = workspace.joinpath("rank_0", "run.log").read_text().splitlines() log_results[-1] += f", {lines[0]}" print(log_results[-1]) results_to_csv(workspace.joinpath("context_perf.csv"), context_results) if xpu_config["perf_config"]["perf_decode"]: batch_size_list = xpu_config["perf_config"]["batch_size_list"] seq_len_list = xpu_config["perf_config"]["seq_len_list"] decode_results = {} for batch_size in batch_size_list: if batch_size > max_batch_size: continue decode_results[batch_size] = {} for seq_len in seq_len_list: input_template = update_template("decode", batch_size, seq_len) start_iters = 2 test_iter = 0 duration_ms = 0. while duration_ms < 5000. and test_iter < 100: result = engine.mp_forward(input_template) if start_iters > 0: start_iters -= 1 continue test_iter += 1 duration_ms += result["duration_ms"] avg_duration_ms = round(duration_ms / test_iter, 3) decode_results[batch_size][seq_len] = avg_duration_ms log_results.append(f"decode, batch_size={batch_size}, seq_len={seq_len}, avg_duration_ms={avg_duration_ms}") if log_model: log_path = workspace.joinpath("rank_0", "run.log") if log_path.exists(): lines = workspace.joinpath("rank_0", "run.log").read_text().splitlines() log_results[-1] += f", {lines[0]}" print(log_results[-1]) results_to_csv(workspace.joinpath("decode_perf.csv"), decode_results) # release subprocess resources manually to enable auto garbage collection engine.clean_subprocess() with open(workspace.joinpath("output.txt"), "w") as f: for log in log_results: f.write(log + "\n") print("\n".join(log_results)) if __name__ == '__main__': args = parse_args() setup_logger(args.log_level) hardware_type = args.hardware_type model_config_path = pathlib.Path(args.model_config) if not model_config_path.is_absolute(): model_config_path = CUR_DIR / model_config_path if not model_config_path.exists(): logger.error(f"model_config_path not exist: {model_config_path}") sys.exit(-1) model_name = model_config_path.stem.split("-") model_name = "-".join([model_name[0], model_name[-1]]) tp_size = args.tp_size max_batch_size = args.max_batch_size model_version = args.model_version if args.workspace is not None: workspace = pathlib.Path(args.workspace) if not workspace.is_absolute(): workspace = CUR_DIR / workspace else: workspace = INFER_ROOT_DIR.joinpath( "llm_perf", "reports", hardware_type, model_config_path.stem, "bench_model" ) if not workspace.exists(): workspace.mkdir(parents=True, exist_ok=True) logger.info(f"hardware_type: {hardware_type}") logger.info(f"model_config: {model_config_path}") logger.info(f"infer_dtype: {args.infer_dtype}") logger.info(f"tp_size: {tp_size}") logger.info(f"max_batch_size: {max_batch_size}") logger.info(f"model_version: {model_version}") logger.info(f"workspace: {workspace}") print("\n\n") xpu_config = { "hardware_type": hardware_type, "sku_name": "", "model_name": model_name, "deploy_config": { "impl_framework": "bytemlperf", "impl_version": model_version, "infer_dtype": args.infer_dtype, "tp_size": tp_size, "max_batch_size": max_batch_size, }, "model_config": json.loads(model_config_path.read_text()), "perf_config": { "perf_context": True, "perf_decode": True, "batch_size_list": [1, 4, 8, 16, 24, 32], "seq_len_list": [1024, 2048, 4096, 8192] }, "feature_config": { "random_seed": args.random_seed, "log_model": args.log_model, }, "tp_size": tp_size, "max_batch_size": max_batch_size, } perf_engine(xpu_config, workspace)