mirror of
https://git.datalinker.icu/vllm-project/vllm.git
synced 2026-04-06 17:27:08 +08:00
233 lines
8.0 KiB
Python
233 lines
8.0 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Benchmark the latency of processing a single batch of requests."""
|
|
|
|
import argparse
|
|
import dataclasses
|
|
import json
|
|
import os
|
|
import random
|
|
import time
|
|
from pathlib import Path
|
|
from typing import Any, Optional, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
from benchmark_utils import (convert_to_pytorch_benchmark_format, get_requests,
|
|
validate_dataset, write_to_json)
|
|
from tqdm import tqdm
|
|
from transformers import AutoTokenizer
|
|
|
|
from vllm import LLM, SamplingParams
|
|
from vllm.engine.arg_utils import EngineArgs
|
|
from vllm.inputs import TextPrompt, TokensPrompt
|
|
from vllm.sampling_params import BeamSearchParams
|
|
from vllm.utils import FlexibleArgumentParser
|
|
|
|
|
|
def save_to_pytorch_benchmark_format(args: argparse.Namespace,
|
|
results: dict[str, Any]) -> None:
|
|
pt_records = convert_to_pytorch_benchmark_format(
|
|
args=args,
|
|
metrics={"latency": results["latencies"]},
|
|
extra_info={k: results[k]
|
|
for k in ["avg_latency", "percentiles"]})
|
|
if pt_records:
|
|
pt_file = f"{os.path.splitext(args.output_json)[0]}.pytorch.json"
|
|
write_to_json(pt_file, pt_records)
|
|
|
|
|
|
def main(args: argparse.Namespace):
|
|
print(args)
|
|
|
|
engine_args = EngineArgs.from_cli_args(args)
|
|
|
|
# NOTE(woosuk): If the request cannot be processed in a single batch,
|
|
# the engine will automatically process the request in multiple batches.
|
|
llm = LLM(**dataclasses.asdict(engine_args))
|
|
assert llm.llm_engine.model_config.max_model_len >= (
|
|
args.input_len +
|
|
args.output_len), ("Please ensure that max_model_len is greater than"
|
|
" the sum of input_len and output_len.")
|
|
|
|
sampling_params = SamplingParams(
|
|
n=args.n,
|
|
temperature=1.0,
|
|
top_p=1.0,
|
|
ignore_eos=True,
|
|
max_tokens=args.output_len,
|
|
detokenize=not args.disable_detokenize,
|
|
)
|
|
print(sampling_params)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(
|
|
args.tokenizer, trust_remote_code=args.trust_remote_code)
|
|
requests = get_requests(args.batch_size, args, tokenizer)
|
|
prompts: list[Union[TextPrompt, TokensPrompt]] = []
|
|
for request in requests:
|
|
prompts.append(
|
|
TokensPrompt(prompt_token_ids=request.prompt["prompt_token_ids"],
|
|
multi_modal_data=request.multi_modal_data)
|
|
if "prompt_token_ids" in request.prompt else \
|
|
TextPrompt(prompt=request.prompt,
|
|
multi_modal_data=request.multi_modal_data))
|
|
|
|
def llm_generate():
|
|
if not args.use_beam_search:
|
|
llm.generate(prompts,
|
|
sampling_params=sampling_params,
|
|
use_tqdm=False)
|
|
else:
|
|
llm.beam_search(
|
|
prompts,
|
|
BeamSearchParams(
|
|
beam_width=args.n,
|
|
max_tokens=args.output_len,
|
|
ignore_eos=True,
|
|
),
|
|
)
|
|
|
|
def run_to_completion(profile_dir: Optional[str] = None):
|
|
if profile_dir:
|
|
with torch.profiler.profile(
|
|
activities=[
|
|
torch.profiler.ProfilerActivity.CPU,
|
|
torch.profiler.ProfilerActivity.CUDA,
|
|
],
|
|
on_trace_ready=torch.profiler.tensorboard_trace_handler(
|
|
str(profile_dir)),
|
|
) as p:
|
|
llm_generate()
|
|
print(p.key_averages().table(sort_by="self_cuda_time_total"))
|
|
else:
|
|
start_time = time.perf_counter()
|
|
llm_generate()
|
|
end_time = time.perf_counter()
|
|
latency = end_time - start_time
|
|
return latency
|
|
|
|
print("Warming up...")
|
|
for _ in tqdm(range(args.num_iters_warmup), desc="Warmup iterations"):
|
|
run_to_completion(profile_dir=None)
|
|
|
|
if args.profile:
|
|
profile_dir = args.profile_result_dir
|
|
if not profile_dir:
|
|
profile_dir = (Path(".") / "vllm_benchmark_result" /
|
|
f"latency_result_{time.time()}")
|
|
print(f"Profiling (results will be saved to '{profile_dir}')...")
|
|
run_to_completion(profile_dir=profile_dir)
|
|
return
|
|
|
|
# Benchmark.
|
|
latencies = []
|
|
for _ in tqdm(range(args.num_iters), desc="Profiling iterations"):
|
|
latencies.append(run_to_completion(profile_dir=None))
|
|
latencies = np.array(latencies)
|
|
percentages = [10, 25, 50, 75, 90, 99]
|
|
percentiles = np.percentile(latencies, percentages)
|
|
print(f"Avg latency: {np.mean(latencies)} seconds")
|
|
for percentage, percentile in zip(percentages, percentiles):
|
|
print(f"{percentage}% percentile latency: {percentile} seconds")
|
|
|
|
# Output JSON results if specified
|
|
if args.output_json:
|
|
results = {
|
|
"avg_latency": np.mean(latencies),
|
|
"latencies": latencies.tolist(),
|
|
"percentiles": dict(zip(percentages, percentiles.tolist())),
|
|
}
|
|
with open(args.output_json, "w") as f:
|
|
json.dump(results, f, indent=4)
|
|
save_to_pytorch_benchmark_format(args, results)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = FlexibleArgumentParser(
|
|
description="Benchmark the latency of processing a single batch of "
|
|
"requests till completion.")
|
|
parser.add_argument("--input-len", type=int, default=32)
|
|
parser.add_argument("--output-len", type=int, default=128)
|
|
parser.add_argument("--batch-size", type=int, default=8)
|
|
parser.add_argument(
|
|
"--n",
|
|
type=int,
|
|
default=1,
|
|
help="Number of generated sequences per prompt.",
|
|
)
|
|
parser.add_argument("--use-beam-search", action="store_true")
|
|
parser.add_argument(
|
|
"--num-iters-warmup",
|
|
type=int,
|
|
default=10,
|
|
help="Number of iterations to run for warmup.",
|
|
)
|
|
parser.add_argument("--num-iters",
|
|
type=int,
|
|
default=30,
|
|
help="Number of iterations to run.")
|
|
parser.add_argument(
|
|
"--profile",
|
|
action="store_true",
|
|
help="profile the generation process of a single batch",
|
|
)
|
|
parser.add_argument(
|
|
"--profile-result-dir",
|
|
type=str,
|
|
default=None,
|
|
help=("path to save the pytorch profiler output. Can be visualized "
|
|
"with ui.perfetto.dev or Tensorboard."),
|
|
)
|
|
parser.add_argument(
|
|
"--output-json",
|
|
type=str,
|
|
default=None,
|
|
help="Path to save the latency results in JSON format.",
|
|
)
|
|
parser.add_argument(
|
|
"--disable-detokenize",
|
|
action="store_true",
|
|
help=("Do not detokenize responses (i.e. do not include "
|
|
"detokenization time in the latency measurement)"),
|
|
)
|
|
parser.add_argument(
|
|
"--dataset-name",
|
|
type=str,
|
|
choices=["sharegpt", "random", "sonnet", "burstgpt", "hf"],
|
|
help="Name of the dataset to benchmark on.",
|
|
default="sharegpt")
|
|
# random dataset
|
|
parser.add_argument(
|
|
"--random-range-ratio",
|
|
type=float,
|
|
default=None,
|
|
help="Range of sampled ratio of input/output length, "
|
|
"used only for RandomDataSet.",
|
|
)
|
|
parser.add_argument("--dataset-path",
|
|
type=str,
|
|
default=None,
|
|
help="Path to the dataset")
|
|
|
|
# LoRA
|
|
parser.add_argument(
|
|
"--lora-path",
|
|
type=str,
|
|
default=None,
|
|
help="Path to the lora adapters to use. This can be an absolute path, "
|
|
"a relative path, or a Hugging Face model identifier.")
|
|
|
|
parser.add_argument("--prefix-len",
|
|
type=int,
|
|
default=None,
|
|
help="Number of prefix tokens per request."
|
|
"This is for the RandomDataset and SonnetDataset")
|
|
|
|
parser = EngineArgs.add_cli_args(parser)
|
|
args = parser.parse_args()
|
|
if args.tokenizer is None:
|
|
args.tokenizer = args.model
|
|
args.backend = "vllm"
|
|
validate_dataset(args)
|
|
random.seed(0)
|
|
main(args)
|