404 lines
14 KiB
Python
404 lines
14 KiB
Python
"""MLC LLM benchmark main entrance"""
|
|
|
|
import functools
|
|
import json
|
|
import random
|
|
from typing import Any, Dict, List, Optional, Tuple # noqa: UP035
|
|
|
|
import numpy as np
|
|
import requests
|
|
from transformers import AutoTokenizer
|
|
|
|
import mlc_llm
|
|
from mlc_llm.bench.api_endpoint import SUPPORTED_BACKENDS, create_api_endpoint
|
|
from mlc_llm.bench.dataset import SUPPORTED_DATASET, Dataset, create_dataset
|
|
from mlc_llm.bench.request_processor import (
|
|
MetricAnalyzer,
|
|
RequestProcessor,
|
|
create_pipelines,
|
|
)
|
|
from mlc_llm.bench.request_record import (
|
|
RequestRecord,
|
|
convert_reports_to_df,
|
|
generate_metrics_summary,
|
|
pretty_print_report,
|
|
)
|
|
from mlc_llm.cli.serve import EngineConfigOverride
|
|
from mlc_llm.serve import EngineConfig
|
|
from mlc_llm.support import argparse, logging
|
|
|
|
logging.enable_logging()
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _parse_num_concurrent_requests(num_str: Optional[str]) -> Optional[List[int]]: # noqa: UP006
|
|
if num_str is None:
|
|
return None
|
|
numbers = num_str.split(",")
|
|
if any(not number.isdigit() for number in numbers):
|
|
raise ValueError(f"Unrecognized num_concurrent_requests list: {numbers}")
|
|
return list(int(number) for number in numbers)
|
|
|
|
|
|
def _parse_request_rate(request_rate_str: Optional[str]) -> Optional[List[np.float32]]: # noqa: UP006
|
|
if request_rate_str is None:
|
|
return None
|
|
request_rates = request_rate_str.split(",")
|
|
results = []
|
|
for rate_str in request_rates:
|
|
request_rate = float(rate_str)
|
|
if request_rate <= 0:
|
|
raise ValueError(f"Invalid request rate {request_rate}")
|
|
results.append(np.float32(request_rate))
|
|
return results
|
|
|
|
|
|
def _parse_mlc_engine_config(config_str: Optional[str]) -> EngineConfig:
|
|
if config_str is None:
|
|
return None
|
|
engine_config_override = EngineConfigOverride.from_str(config_str)
|
|
return EngineConfig(
|
|
tensor_parallel_shards=engine_config_override.tensor_parallel_shards,
|
|
max_num_sequence=engine_config_override.max_num_sequence,
|
|
max_total_sequence_length=engine_config_override.max_total_seq_length,
|
|
prefill_chunk_size=engine_config_override.prefill_chunk_size,
|
|
sliding_window_size=engine_config_override.sliding_window_size,
|
|
attention_sink_size=engine_config_override.attention_sink_size,
|
|
max_history_size=engine_config_override.max_history_size,
|
|
gpu_memory_utilization=engine_config_override.gpu_memory_utilization,
|
|
spec_draft_length=engine_config_override.spec_draft_length,
|
|
prefill_mode=engine_config_override.prefill_mode,
|
|
prefix_cache_max_num_recycling_seqs=engine_config_override.prefix_cache_max_num_recycling_seqs,
|
|
prefix_cache_mode=engine_config_override.prefix_cache_mode,
|
|
)
|
|
|
|
|
|
def _launch_mlc_server(args: argparse.argparse.Namespace):
|
|
return mlc_llm.serve.PopenServer(
|
|
model=args.tokenizer,
|
|
mode="server",
|
|
model_lib=args.mlc_model_lib,
|
|
enable_tracing=False,
|
|
host=args.host,
|
|
port=args.port,
|
|
engine_config=args.mlc_engine_config,
|
|
)
|
|
|
|
|
|
def run_pipeline(
|
|
pipeline: RequestProcessor,
|
|
dataset: Dataset,
|
|
tokenizer: AutoTokenizer,
|
|
args: argparse.argparse.Namespace,
|
|
) -> Tuple[Dict[str, Any], List[RequestRecord]]: # noqa: UP006
|
|
"""Run the pipeline with the given dataset and args. Return the benchmark report dict."""
|
|
random.seed(args.seed)
|
|
np.random.seed(args.seed)
|
|
request_records = dataset.generate_request_records(
|
|
args.input_len,
|
|
args.output_len,
|
|
args.input_len_std,
|
|
args.output_len_std,
|
|
)
|
|
request_records = pipeline(request_records)
|
|
num_total_requests = (
|
|
args.num_requests if not args.per_gpu_workload else args.num_requests * args.num_gpus
|
|
)
|
|
assert len(request_records) == num_total_requests
|
|
sorted_requests: List[RequestRecord] = [None] * num_total_requests # noqa: UP006
|
|
for request_record in request_records:
|
|
assert request_record.request_id is not None
|
|
assert sorted_requests[request_record.request_id] is None
|
|
sorted_requests[request_record.request_id] = request_record
|
|
|
|
request_records = MetricAnalyzer(tokenizer)(request_records)
|
|
report = generate_metrics_summary(request_records, num_total_requests, args.num_gpus)
|
|
return report, sorted_requests
|
|
|
|
|
|
def query_mlc_server_metrics(host: str, port: int):
|
|
"""Try to get the MLC server metrics whenever it exists."""
|
|
try:
|
|
r = requests.post(f"http://{host}:{port}/debug/dump_engine_metrics", json={}, timeout=10)
|
|
if r.status_code == 200:
|
|
print(f"MLC server metrics: {r.json()}")
|
|
except Exception:
|
|
pass
|
|
|
|
|
|
def main(args: argparse.argparse.Namespace):
|
|
"""Main benchmark entrance."""
|
|
mlc_server = None
|
|
if args.mlc_model_lib:
|
|
mlc_server = _launch_mlc_server(args)
|
|
if args.num_requests <= 0:
|
|
raise ValueError("Number of requests to benchmark must be positive.")
|
|
|
|
def _main():
|
|
tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
|
|
dataset = create_dataset(args, tokenizer)
|
|
f_create_api_endpoint = functools.partial(create_api_endpoint, args)
|
|
pipelines = create_pipelines(args, f_create_api_endpoint, dataset)
|
|
reports = []
|
|
alltime_records = {}
|
|
for i, pipeline in enumerate(pipelines):
|
|
report, request_records = run_pipeline(pipeline, dataset, tokenizer, args)
|
|
exec_feature = (
|
|
json.dumps(report["exec_feature"])
|
|
if report["exec_feature"] is not None
|
|
else f"pipeline{i}"
|
|
)
|
|
alltime_records[exec_feature] = [
|
|
request_record.model_dump() for request_record in request_records
|
|
]
|
|
reports.append(report)
|
|
pretty_print_report(report)
|
|
query_mlc_server_metrics(args.host, args.port)
|
|
|
|
# Construct data frame
|
|
df = convert_reports_to_df(reports)
|
|
print(df)
|
|
df.to_csv(args.output, index=False)
|
|
logger.info("Benchmark results dumped to file %s", args.output)
|
|
if args.debug_dump:
|
|
debug_dump_filepath = (
|
|
args.output[:-4] if args.output.endswith(".csv") else args.output
|
|
) + "_debug_dump.log"
|
|
with open(debug_dump_filepath, "w", encoding="utf-8") as file:
|
|
json.dump(alltime_records, file, indent=4)
|
|
logger.info("Debug log dumped to file %s", debug_dump_filepath)
|
|
|
|
if mlc_server is not None:
|
|
with mlc_server:
|
|
_main()
|
|
else:
|
|
_main()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser("MLC LLM benchmark")
|
|
|
|
parser.add_argument(
|
|
"--dataset",
|
|
type=str,
|
|
choices=SUPPORTED_DATASET,
|
|
help=f"The benchmark dataset kind. Supporting {SUPPORTED_DATASET}",
|
|
)
|
|
parser.add_argument(
|
|
"--dataset-path",
|
|
type=str,
|
|
help="The dataset file path.",
|
|
)
|
|
parser.add_argument(
|
|
"--api-endpoint",
|
|
type=str,
|
|
choices=SUPPORTED_BACKENDS,
|
|
default="openai",
|
|
help="The API endpoint API for benchmarking.",
|
|
)
|
|
parser.add_argument(
|
|
"--tokenizer",
|
|
type=str,
|
|
required=True,
|
|
help="The path of the tokenizer directory.",
|
|
)
|
|
parser.add_argument(
|
|
"--num-gpus",
|
|
type=int,
|
|
required=True,
|
|
help="The number of GPUs used by the server. "
|
|
"We need this to better analyze the throughput per GPU.",
|
|
)
|
|
parser.add_argument(
|
|
"--num-requests",
|
|
type=int,
|
|
required=True,
|
|
help="The number of requests for benchmark.",
|
|
)
|
|
parser.add_argument(
|
|
"--num-warmup-requests",
|
|
type=int,
|
|
help="The number of requests for warmup. "
|
|
"It is optional when fixing the number of concurrent requests, and is required otherwise.",
|
|
)
|
|
parser.add_argument(
|
|
"--per-gpu-workload",
|
|
default=False,
|
|
action="store_true",
|
|
help='When set to True, the specified "num_concurrent_requests"/"request_rate" '
|
|
"denote the workload **per GPU**, which means that the real values of "
|
|
'"num_concurrent_requests"/"request_rate" used in benchmark'
|
|
'will be multiplied by "num_gpus".',
|
|
)
|
|
parser.add_argument(
|
|
"--num-concurrent-requests",
|
|
type=_parse_num_concurrent_requests,
|
|
help="The number(s) of concurrent requests to benchmark. "
|
|
'It can be either one integer or a list of integer separated by commas(","). '
|
|
"When specified, for each integer, the benchmark keeps these many consistent "
|
|
"number of concurrently running requests.",
|
|
)
|
|
parser.add_argument(
|
|
"--request-rate",
|
|
type=_parse_request_rate,
|
|
help="The request rate(s) denoting the number of new requests each second. "
|
|
'It can be either one float number (or "inf") or a list of numbers separated '
|
|
'by commas(","). '
|
|
"When specified, the benchmark sends these many new requests each second. "
|
|
'If it is "inf", all requests will be sent together at once.',
|
|
)
|
|
parser.add_argument(
|
|
"--replay-timestamp-scale",
|
|
type=float,
|
|
help="The timestamp scale when replaying the timestamps in a dataset. "
|
|
'The dataset replay mode is enabled when neither "--num-concurrent-requests" and '
|
|
'"--request-rate" is specified. '
|
|
"The scale is 1 by default in the replay mode.",
|
|
)
|
|
parser.add_argument(
|
|
"--input-len",
|
|
type=int,
|
|
help="The benchmark request average input length. Default to None, "
|
|
"which means the request input length depends on the dataset being used.",
|
|
)
|
|
parser.add_argument(
|
|
"--input-len-std",
|
|
type=float,
|
|
default=0,
|
|
help="The benchmark request input length standard deviation. Default to 0.",
|
|
)
|
|
parser.add_argument(
|
|
"--output-len",
|
|
type=int,
|
|
help="The benchmark request average output length. Default to None, "
|
|
"which means the request output length depends on the dataset being used.",
|
|
)
|
|
parser.add_argument(
|
|
"--output-len-std",
|
|
type=float,
|
|
default=0,
|
|
help="The benchmark request output length standard deviation. Default to 0.",
|
|
)
|
|
parser.add_argument(
|
|
"--stream",
|
|
type=bool,
|
|
default=True,
|
|
help="Whether to benchmark stream responses. "
|
|
"When not enabled, metrics such as time-to-first-token (TTFT) will not be available. "
|
|
"Default to True.",
|
|
)
|
|
parser.add_argument(
|
|
# NOTE: The current implementation of server metrics still has some issues that need fixes,
|
|
# which makes it not work to include server metrics.
|
|
"--include-server-metrics",
|
|
action="store_true",
|
|
help="Whether to also benchmark the server side request metrics. "
|
|
"This option is only available when benchmarking MLC server.",
|
|
)
|
|
parser.add_argument(
|
|
"--host",
|
|
type=str,
|
|
required=True,
|
|
help="The host address of the backend API.",
|
|
)
|
|
parser.add_argument(
|
|
"--port",
|
|
type=int,
|
|
required=True,
|
|
help="The port of the backend API.",
|
|
)
|
|
parser.add_argument(
|
|
"--timeout",
|
|
type=float,
|
|
default=3 * 60 * 60,
|
|
help="The timeout limit of each request.",
|
|
)
|
|
parser.add_argument(
|
|
"--seed",
|
|
type=int,
|
|
default=0,
|
|
help="The random number seed. Default to 0.",
|
|
)
|
|
parser.add_argument(
|
|
"--temperature",
|
|
type=float,
|
|
default=1.0,
|
|
help="The temperature value for logit adjustment. Default to 1.",
|
|
)
|
|
parser.add_argument(
|
|
"--top-p",
|
|
type=float,
|
|
default=1.0,
|
|
help="The top-p value for sampling. Default to 1.",
|
|
)
|
|
parser.add_argument(
|
|
"--ignore-eos",
|
|
default=False,
|
|
action="store_true",
|
|
help='Whether to set the "ignore_eos" field.',
|
|
)
|
|
parser.add_argument(
|
|
"--apply-chat-template",
|
|
default=False,
|
|
action="store_true",
|
|
help="Whether to apply chat template to the request input text. "
|
|
'It is not supported when "--input-len" is specified.',
|
|
)
|
|
parser.add_argument(
|
|
"--num-process-workers",
|
|
type=int,
|
|
help="The number of parallel process workers to send the requests.",
|
|
)
|
|
parser.add_argument(
|
|
"--disable-tqdm",
|
|
action="store_true",
|
|
help="Whether to disable showing progress bar with tqdm during benchmarking.",
|
|
)
|
|
parser.add_argument(
|
|
"--max-schedule-gap",
|
|
type=float,
|
|
default=0.5,
|
|
help="The maximum allowed delay between the scheduled time in seconds.",
|
|
)
|
|
parser.add_argument(
|
|
"--mlc-model-lib",
|
|
type=str,
|
|
help="The model lib path when benchmarking MLC serve. "
|
|
"When specified, the server is automatic launched and no external server launch is needed.",
|
|
)
|
|
parser.add_argument(
|
|
"--mlc-engine-config",
|
|
type=_parse_mlc_engine_config,
|
|
help="The engine config used when launch MLC server.",
|
|
)
|
|
parser.add_argument(
|
|
"--cuda-profile",
|
|
default=False,
|
|
action="store_true",
|
|
help="Whether to enable cuda profile on server. "
|
|
"The --mlc-model-lib path should be provided when enabling this option.",
|
|
)
|
|
parser.add_argument(
|
|
"--debug-dump",
|
|
default=False,
|
|
action="store_true",
|
|
help="Whether to dump all request record raw data to file.",
|
|
)
|
|
parser.add_argument(
|
|
"--multi-round",
|
|
default=False,
|
|
action="store_true",
|
|
help="Whether to chat like multi round conversion with history log each request. "
|
|
"Only enabled when benchmarked with fixed concurrent request mode."
|
|
"The --num-concurrent-requests should be provided when enabling this option.",
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
"-o",
|
|
type=str,
|
|
default="mlc_benchmark.csv",
|
|
help="The path of the output file where to dump the benchmark results.",
|
|
)
|
|
|
|
main(parser.parse_args())
|