saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-25 19:48:11 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add LLaMA end-to-end benchmarking (#19985)

Browse source 
### Description

This PR adds a benchmarking script to measure end-to-end performance and
saves the results in a CSV file.

### Motivation and Context

With this PR, end-to-end performance can be easily measured for many
large-language models such as LLaMA-2. The performance numbers for
LLaMA-2 are located
[here](https://github.com/microsoft/onnxruntime-inference-examples/tree/main/python/models/llama).
This commit is contained in:
kunal-vaishnavi 2024-03-21 18:59:05 -07:00 committed by GitHub
parent eab35c20fc
commit 6238e9c0af
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 957 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

132
onnxruntime/python/tools/transformers/models/llama/README.md
View file

@ -1,7 +1,14 @@
# Contents
- [LLaMA-2](#llama-2)
- [Prerequisites](#prerequisites)
- [Exporting LLaMA-2](#exporting-llama-2)
- [Examples of Exporting LLaMA-2](#examples-of-exporting-llama-2)
- [Parity Checking LLaMA-2](#parity-checking-llama-2)
- [Benchmarking LLaMA-2](#benchmark-llama-2)
- [Variants](#variants)
- [Benchmark All](#benchmark-all)
- [Benchmark E2E](#benchmark-e2e)
- [E2E Inference with LLaMA-2](#e2e-inference-with-llama-2)
- [Mistral](#mistral)
- [Exporting Mistral](#exporting-mistral)
- [Optimizing and Quantizing Mistral](#optimizing-and-quantizing-mistral)
@ -229,6 +236,55 @@ $ ./build.sh --config Release --use_cuda --cuda_home /usr/local/cuda-12.2 --cudn
$ CUDA_VISIBLE_DEVICES=0,1,2,3 bash convert_70b_model.sh 4 -m meta-llama/Llama-2-70b-hf --output llama2-70b-distributed --precision fp16 --execution_provider cuda --use_gqa
```
## Parity Checking LLaMA-2
Here are some examples of how you can use the parity checker to verify your LLaMA-2 ONNX model.
1. Merged ONNX model, FP32 CPU
```
CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.llama_parity \
--model_name meta-llama/Llama-2-7b-hf \
--onnx_model_path ./llama2-7b/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32.onnx \
--merged \
--execution_provider cpu \
--precision fp32 \
--cache_dir ./model_cache \
```
2. Merged ONNX model, FP32 CUDA
```
CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.llama_parity \
--model_name meta-llama/Llama-2-7b-hf \
--onnx_model_path ./llama2-7b/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32.onnx \
--merged \
--execution_provider cuda \
--precision fp32 \
--cache_dir ./model_cache \
```
3. Merged ONNX model, FP16 CUDA
```
CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.llama_parity \
--model_name meta-llama/Llama-2-7b-hf \
--onnx_model_path ./llama2-7b/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32.onnx \
--merged \
--execution_provider cuda \
--precision fp16 \
--cache_dir ./model_cache \
```
4. Merged ONNX model, FP16 CUDA with GroupQueryAttention
```
CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.llama_parity \
--model_name meta-llama/Llama-2-7b-hf \
--onnx_model_path ./llama2-7b/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32.onnx \
--merged \
--use_gqa \
--execution_provider cuda \
--precision fp16 \
--cache_dir ./model_cache \
```
## Benchmark LLaMA-2
Here are some examples of how you can benchmark LLaMA-2.
@ -240,6 +296,7 @@ Here are some examples of how you can benchmark LLaMA-2.
CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark \
--benchmark-type hf-pt-eager \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--precision fp32 \
--batch-sizes "1 2" \
--sequence-lengths "8 16" \
@ -252,6 +309,7 @@ CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark \
CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark \
--benchmark-type hf-pt-compile \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--precision fp16 \
--batch-sizes "1 2" \
--sequence-lengths "8 16" \
@ -265,6 +323,7 @@ CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark \
--benchmark-type hf-ort \
--hf-ort-dir-path ./Llama-2-7b-hf-onnx/ \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--precision fp32 \
--batch-sizes "1 2" \
--sequence-lengths "8 16" \
@ -278,6 +337,7 @@ CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark \
--benchmark-type hf-ort \
--hf-ort-dir-path ./Llama-2-7b-hf-onnx/ \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--precision fp16 \
--batch-sizes "1 2" \
--sequence-lengths "8 16" \
@ -291,6 +351,7 @@ CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark \
--benchmark-type ort-msft \
--ort-model-path ./llama-2-onnx/7B_float32/ONNX/LlamaV2_7B_float32.onnx \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--precision fp32 \
--batch-sizes "1 2" \
--sequence-lengths "8 16" \
@ -303,6 +364,7 @@ CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark \
--benchmark-type ort-msft \
--ort-model-path ./llama-2-onnx/7B_float16/ONNX/LlamaV2_7B_float16.onnx \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--precision fp16 \
--batch-sizes "1 2" \
--sequence-lengths "8 16" \
@ -315,6 +377,7 @@ CUDA_VISIBLE_DEVICES=1 python3 -m models.llama.benchmark \
--benchmark-type ort-convert-to-onnx \
--ort-model-path ./llama2-7b/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32.onnx \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--precision fp32 \
--batch-sizes "1 2" \
--sequence-lengths "8 16" \
@ -327,6 +390,7 @@ CUDA_VISIBLE_DEVICES=4 python3 -m models.llama.benchmark \
--benchmark-type ort-convert-to-onnx \
--ort-model-path ./llama2-7b/rank_0_Llama-2-7b-hf_decoder_merged_model_fp16.onnx \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--precision fp16 \
--batch-sizes "1 2" \
--sequence-lengths "8 16" \
@ -339,6 +403,7 @@ CUDA_VISIBLE_DEVICES=4,5,6,7 bash benchmark_70b_model.sh 4 \
--benchmark-type ort-convert-to-onnx \
--ort-model-path ./llama2-70b-dis/rank_{}_Llama-2-70b-hf_decoder_merged_model_fp16.onnx \
--model-name meta-llama/Llama-2-70b-hf \
--cache-dir ./model_cache \
--precision fp16 \
--device cuda \
--warmup-runs 5 \
@ -357,6 +422,7 @@ CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark_all \
--ort-convert-to-onnx-model-path ./llama2-7b-fp16/Llama-2-7b-hf_decoder_merged_model_fp16.onnx \
--ort-msft-model-path ./llama-2-onnx/7B_float16/ONNX/LlamaV2_7B_float16.onnx \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--precision fp16 \
--batch-sizes "1 2" \
--sequence-lengths "8 16" \
@ -366,6 +432,72 @@ CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark_all \
--timeout 60 # number of minutes before moving to the next benchmark
```
### Benchmark E2E
You can use `benchmark_e2e.py` to benchmark the full end-to-end scenario and automatically store the results in a CSV file. This tool uses `argmax` for sampling to standardize the benchmarking process.
1. PyTorch without `torch.compile`, FP32
```
CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark_e2e \
--benchmark-type pt-eager \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--prompts-file ./models/llama/prompts.json \
--precision fp32 \
--batch-sizes "1 2" \
--prompt-lengths "16 64" \
--device cpu \
--auth
```
2. PyTorch with `torch.compile`, FP16
```
CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark_e2e \
--benchmark-type pt-compile \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--prompts-file ./models/llama/prompts.json \
--precision fp16 \
--batch-sizes "1 2" \
--prompt-lengths "16 64" \
--device cuda \
--auth
```
3. ONNX Runtime with `convert_to_onnx`, FP32
```
CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark_e2e \
--benchmark-type ort \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--onnx-model-path ./llama2-7b/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32.onnx \
--prompts-file ./models/llama/prompts.json \
--precision fp32 \
--batch-sizes "1 2" \
--prompt-lengths "16 64" \
--device cpu \
--auth
```
4. ONNX Runtime with `convert_to_onnx`, FP16
```
CUDA_VISIBLE_DEVICES=0 python3 -m models.llama.benchmark_e2e \
--benchmark-type ort \
--model-name meta-llama/Llama-2-7b-hf \
--cache-dir ./model_cache \
--onnx-model-path ./llama2-7b/rank_0_Llama-2-7b-hf_decoder_merged_model_fp32.onnx \
--prompts-file ./models/llama/prompts.json \
--precision fp16 \
--batch-sizes "1 2" \
--prompt-lengths "16 64" \
--device cuda \
--use_buffer_share \
--auth
```
## E2E Inference with LLaMA-2
For end-to-end inference, please visit the [ONNX Runtime Inference Examples folder](https://github.com/microsoft/onnxruntime-inference-examples/tree/main/python/models/llama) for a step-by-step walkthrough, code examples, and performance metrics.
# Mistral
## Introduction

38
onnxruntime/python/tools/transformers/models/llama/benchmark.py
View file

@ -1,3 +1,8 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
import argparse
import datetime
import gc
@ -14,11 +19,12 @@ import torch
from benchmark_helper import measure_memory, setup_logger
from dist_settings import get_rank, get_size
from llama_inputs import (
add_io_bindings,
add_io_bindings_as_ortvalues,
get_merged_sample_with_past_kv_inputs,
get_msft_sample_inputs,
get_sample_inputs,
get_sample_with_past_kv_inputs,
verify_ort_inputs,
)
from optimum.onnxruntime import ORTModelForCausalLM
from torch.profiler import ProfilerActivity, profile, record_function
@ -199,6 +205,7 @@ def get_model(args: argparse.Namespace):
torch_dtype=torch.float16 if args.use_fp16 else torch.float32,
use_auth_token=args.auth,
use_cache=True,
cache_dir=args.cache_dir,
).to(args.target_device)
end_time = time.time()
@ -444,24 +451,12 @@ def run_hf_inference(args, init_inputs, iter_inputs, model):
def run_ort_inference(args, init_inputs, iter_inputs, model):
def prepare_ort_inputs(inputs, kv_cache_ortvalues):
# Check that all model inputs will be provided
model_inputs = set(map(lambda model_input: model_input.name, model.get_inputs()))
user_inputs = set(inputs.keys())
missing_inputs = model_inputs - user_inputs
if len(missing_inputs):
logger.error(f"The following model inputs are missing: {missing_inputs}")
raise Exception("There are missing inputs to the model. Please add them and try again.")
# Remove unnecessary inputs from model inputs
unnecessary_inputs = user_inputs - model_inputs
if len(unnecessary_inputs):
for unnecessary_input in unnecessary_inputs:
logger.info(f"Removing unnecessary input '{unnecessary_input}' from user provided inputs")
del inputs[unnecessary_input]
# Verify model inputs
inputs = verify_ort_inputs(model, inputs)
# Add IO bindings for non-CPU execution providers
if args.device != "cpu":
io_binding, kv_cache_ortvalues = add_io_bindings(
io_binding, kv_cache_ortvalues = add_io_bindings_as_ortvalues(
model, inputs, args.device, int(args.rank), args.use_gqa, kv_cache_ortvalues
)
setattr(args, "io_binding", io_binding) # noqa: B010
@ -612,6 +607,13 @@ def get_args(rank=0):
parser.add_argument("--pt-num-rows", type=int, default=1000, help="Number of rows for PyTorch profiler to display")
parser.add_argument("--verbose", default=False, action="store_true")
parser.add_argument("--log-folder", type=str, default=os.path.join("."), help="Folder to cache log files")
parser.add_argument(
"--cache-dir",
type=str,
required=True,
default="./model_cache",
help="Cache dir where Hugging Face files are stored",
)
args = parser.parse_args()
@ -662,8 +664,8 @@ def main():
args.rank = rank
args.world_size = world_size
tokenizer = AutoTokenizer.from_pretrained(args.model_name)
config = AutoConfig.from_pretrained(args.model_name)
tokenizer = AutoTokenizer.from_pretrained(args.model_name, cache_dir=args.cache_dir)
config = AutoConfig.from_pretrained(args.model_name, cache_dir=args.cache_dir)
target_device = f"cuda:{args.rank}" if args.device != "cpu" else args.device
use_fp16 = args.precision == "fp16"

22
onnxruntime/python/tools/transformers/models/llama/benchmark_all.py
View file

@ -1,3 +1,8 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
import argparse
import datetime
import json
@ -78,6 +83,13 @@ def get_args():
help="Path to ONNX model from convert_to_onnx",
)
parser.add_argument(
"--cache-dir",
type=str,
default="./model_cache",
help="Cache dir where Hugging Face files are stored",
)
parser.add_argument(
"--model-name",
type=str,
@ -332,6 +344,8 @@ def main():
str(args.num_runs),
"--log-folder",
args.log_folder,
"--cache-dir",
args.cache_dir,
"--auth",
]
logger.info("Benchmark PyTorch without torch.compile")
@ -362,6 +376,8 @@ def main():
str(args.num_runs),
"--log-folder",
args.log_folder,
"--cache-dir",
args.cache_dir,
"--auth",
]
logger.info("Benchmark PyTorch with torch.compile")
@ -394,6 +410,8 @@ def main():
str(args.num_runs),
"--log-folder",
args.log_folder,
"--cache-dir",
args.cache_dir,
"--auth",
]
logger.info("Benchmark Optimum + ONNX Runtime")
@ -426,6 +444,8 @@ def main():
str(args.num_runs),
"--log-folder",
args.log_folder,
"--cache-dir",
args.cache_dir,
]
logger.info("Benchmark Microsoft model in ONNX Runtime")
results = benchmark(args, benchmark_cmd, "ort-msft")
@ -457,6 +477,8 @@ def main():
str(args.num_runs),
"--log-folder",
args.log_folder,
"--cache-dir",
args.cache_dir,
]
logger.info("Benchmark convert_to_onnx model in ONNX Runtime")
results = benchmark(args, benchmark_cmd, "onnxruntime")

554
onnxruntime/python/tools/transformers/models/llama/benchmark_e2e.py Normal file
View file

@ -0,0 +1,554 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
# This is an end-to-end benchmarking script for the Hugging Face LLaMA-2 model.
#
# Prerequisites:
# 1) Install `huggingface-cli`:
#
# $ pip install huggingface_hub
#
# 2) Authenticate with Hugging Face's CLI:
#
# $ huggingface-cli login
#
# 3) Accept Meta's license in Hugging Face to access the models at https://huggingface.co/meta-llama/
#
# 4) Install the latest ONNX Runtime version
#
# $ pip install onnxruntime-gpu
from __future__ import annotations
import argparse
import datetime
import gc
import itertools
import json
import logging
import os
import textwrap
import time
import numpy as np
import pandas as pd
import torch
from benchmark_helper import setup_logger
from llama_inputs import add_io_bindings_as_tensors, get_initial_inputs_and_outputs
from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
import onnxruntime as ort
logger = logging.getLogger(__name__)
def get_model(args):
if args.benchmark_type in {"pt-eager", "pt-compile"}:
model = AutoModelForCausalLM.from_pretrained(
args.hf_dir_path if args.hf_dir_path != "" else args.model_name,
cache_dir=args.cache_dir,
torch_dtype=args.torch_dtype,
use_auth_token=args.auth,
use_cache=True,
).to(args.target_device)
model.eval()
if args.benchmark_type == "pt-compile":
model = torch.compile(model)
else:
sess_options = ort.SessionOptions()
ep = (
("CUDAExecutionProvider", {"device_id": args.device_id})
if args.device == "cuda"
else "CPUExecutionProvider"
)
model = ort.InferenceSession(args.onnx_model_path, sess_options=sess_options, providers=[ep])
return model
def run_inference(args, model, runs, inputs, outputs):
if args.benchmark_type == "pt-compile":
with torch.no_grad():
outputs = model(**inputs)
# Synchronize inputs
io_binding = None
if args.benchmark_type in {"pt-eager", "pt-compile"}:
if args.device != "cpu":
torch.cuda.synchronize(args.target_device)
else:
io_binding = add_io_bindings_as_tensors(model, inputs, outputs, args.use_fp16, args.use_buffer_share)
io_binding.synchronize_inputs()
# Run inference
start = time.perf_counter()
for _ in range(runs):
if args.benchmark_type in {"pt-eager", "pt-compile"}:
with torch.no_grad():
outputs = model(**inputs)
if args.device != "cpu":
torch.cuda.synchronize(args.target_device)
else:
model.run_with_iobinding(io_binding)
io_binding.synchronize_outputs()
end = time.perf_counter()
avg = (end - start) / runs
return avg, outputs
def prepare_model_for_inference(args, model, config, tokenizer, prompt_length, prompt):
clear_cache()
inputs, outputs = get_initial_inputs_and_outputs(
config, tokenizer, prompt_length, prompt, args.target_device, args.use_fp16, args.use_buffer_share, args.engine
)
_, outputs = run_inference(args, model, args.warmup_runs, inputs, outputs)
return inputs, outputs
def clear_cache():
gc.collect()
torch.cuda.empty_cache()
def save_results(results, filename, gen_length):
df = pd.DataFrame(
results,
columns=[
"Batch Size",
"Prompt Length",
"Prompt Processing Latency (ms)",
"Prompt Processing Throughput (tps)",
"Sampling Latency (ms)",
"Sampling Throughput (tps)",
"First Token Generated Latency (ms)",
"First Token Generated Throughput (tps)",
f"Average Latency of First {gen_length // 2} Tokens Generated (ms)",
f"Average Throughput of First {gen_length // 2} Tokens Generated (tps)",
f"Average Latency of First {gen_length} Tokens Generated (ms)",
f"Average Throughput of First {gen_length} Tokens Generated (tps)",
"Wall-Clock Latency (s)",
"Wall-Clock Throughput (tps)",
],
)
df.to_csv(filename, index=False)
logger.info(f"Results saved in {filename}!")
def get_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"-bt",
"--benchmark-type",
type=str,
required=True,
choices=["pt-eager", "pt-compile", "ort"],
)
parser.add_argument(
"-m",
"--model-name",
type=str,
required=False,
help="Hugging Face name of model (e.g. 'meta-llama/Llama-2-7b-hf')",
)
parser.add_argument(
"-a",
"--auth",
default=False,
action="store_true",
help="Use Hugging Face authentication token to access model",
)
parser.add_argument(
"-c",
"--cache-dir",
type=str,
default=os.path.join(".", "model_cache"),
help="Path to directory containing all Hugging Face files (e.g. config, tokenizer, PyTorch model). Use when loading model as `AutoModel.from_pretrained(model_name, cache_dir=cache_dir)`.",
)
parser.add_argument(
"--hf-dir-path",
type=str,
default="",
help="Path to directory containing all Hugging Face files (e.g. config, tokenizer, PyTorch model). Use when loading model as `AutoModel.from_pretrained(folder_path)`.",
)
parser.add_argument(
"-o",
"--onnx-model-path",
required=False,
help="Path to ONNX model",
)
parser.add_argument(
"-f",
"--prompts-file",
required=True,
default=os.path.join(".", "models", "llama", "prompts.json"),
help="JSON file containing entries in the format 'prompt length: prompt' where prompt length = tokenized length of prompt",
)
parser.add_argument(
"--use_buffer_share",
default=False,
action="store_true",
help="Use when GroupQueryAttention (GQA) is in ONNX model",
)
parser.add_argument(
"--anomaly-filtering",
default=False,
action="store_true",
help="Use this flag to filter anomaly accelerator times for tokens generated. \
This may give more accurate latency and throughput metrics for tokens generated. \
Wall-clock metrics are still reported with anomaly times though.",
),
parser.add_argument(
"-b",
"--batch-sizes",
default="1 2",
)
parser.add_argument(
"-s",
"--prompt-lengths",
default="32 64 128 256 512",
)
parser.add_argument(
"-p",
"--precision",
required=True,
type=str,
default="fp32",
choices=["int4", "int8", "fp16", "fp32"],
help="Precision for model. For ONNX models, the model's precision should be set before running this script.",
)
parser.add_argument(
"-g",
"--generation-length",
type=int,
default=256,
help="Number of new tokens to generate",
)
parser.add_argument(
"-d",
"--device",
type=str,
default="cuda" if torch.cuda.is_available() else "cpu",
choices=["cpu", "cuda"],
)
parser.add_argument("-id", "--device-id", type=int, default=0)
parser.add_argument("-w", "--warmup-runs", type=int, default=5)
parser.add_argument("-n", "--num-runs", type=int, default=100)
parser.add_argument("--seed", type=int, default=2)
args = parser.parse_args()
# Set seed properties
np.random.seed(args.seed)
torch.manual_seed(args.seed)
# Set runtime properties
if "ort" in args.benchmark_type:
setattr(args, "execution_provider", f"{args.device.upper()}ExecutionProvider") # noqa: B010
if args.execution_provider == "CUDAExecutionProvider":
args.execution_provider = (args.execution_provider, {"device_id": args.device_id})
# Check that paths have been specified for any benchmarking with ORT
if args.benchmark_type == "ort":
assert args.onnx_model_path, "Please specify a path to `--onnx-model-path`"
args.batch_sizes = args.batch_sizes.split(" ")
args.prompt_lengths = args.prompt_lengths.split(" ")
# Use FP32 precision for FP32, INT8, INT4 CPU models, use FP16 precision for FP16 and INT4 GPU models
args.precision = (
"fp32" if args.precision in {"int8", "fp32"} or (args.precision == "int4" and args.device == "cpu") else "fp16"
)
target_device = f"cuda:{args.device_id}" if args.device != "cpu" else args.device
torch_dtype = torch.float16 if args.precision == "fp16" else torch.float32
engine = "ort" if args.benchmark_type == "ort" else "pt"
setattr(args, "target_device", target_device) # noqa: B010
setattr(args, "torch_dtype", torch_dtype) # noqa: B010
setattr(args, "engine", engine) # noqa: B010
setattr(args, "use_fp16", args.precision == "fp16") # noqa: B010
return args
def main():
args = get_args()
setup_logger(False)
logger.info(args.__dict__)
# Get prompts and prompt sizes
size_to_prompt = None
with open(args.prompts_file) as f:
size_to_prompt = json.load(f, object_hook=lambda d: {int(k): v for k, v in d.items()})
# Get config, tokenizer, and model
config = AutoConfig.from_pretrained(
args.hf_dir_path if args.hf_dir_path != "" else args.model_name,
cache_dir=args.cache_dir,
use_auth_token=args.auth,
)
tokenizer = AutoTokenizer.from_pretrained(
args.hf_dir_path if args.hf_dir_path != "" else args.model_name,
cache_dir=args.cache_dir,
use_auth_token=args.auth,
)
model = get_model(args)
all_csv_metrics = []
for batch_size, prompt_length in itertools.product(args.batch_sizes, args.prompt_lengths):
batch_size, prompt_length = int(batch_size), int(prompt_length) # noqa: PLW2901
logger.info(f"Running batch size = {batch_size}, prompt length = {prompt_length}")
clear_cache()
max_length = prompt_length + args.generation_length
if prompt_length not in size_to_prompt:
raise NotImplementedError(
textwrap.dedent(
f"""
A prompt of size {prompt_length} was not found in '{args.prompts_file}'. There are a couple of solutions to fix this.
1) You can change one of the keys in '{args.prompts_file}' to be {prompt_length}.
If {prompt_length} < actual prompt's length, the benchmark E2E tool will repeat the first word in the prompt until {prompt_length} = actual prompt's length.
If {prompt_length} > actual prompt's length, the benchmark E2E tool will automatically trim the actual prompt's length so that {prompt_length} = actual prompt's length.
2) You can add a new key-value entry in '{args.prompts_file}' of the form '{prompt_length}': 'your prompt goes here'.
"""
)
)
prompt = [size_to_prompt[prompt_length]] * batch_size
csv_metrics = [batch_size, prompt_length]
try:
# Measure prompt processing
logger.info("Measuring prompt processing...")
inputs, outputs = prepare_model_for_inference(args, model, config, tokenizer, prompt_length, prompt)
accelerator_prompt_latency_s, outputs = run_inference(args, model, args.num_runs, inputs, outputs)
# Calculate prompt metrics
accelerator_prompt_latency_ms = accelerator_prompt_latency_s * 1000
accelerator_prompt_thrpt = batch_size * (prompt_length / accelerator_prompt_latency_s)
logger.info(f"Average Latency of Prompt Processing: {accelerator_prompt_latency_ms} ms")
logger.info(
f"Average Throughput of Prompt Processing: {batch_size * (prompt_length / accelerator_prompt_latency_s)} tps"
)
csv_metrics.extend([accelerator_prompt_latency_ms, accelerator_prompt_thrpt])
# Measure token generation
logger.info("Measuring token generation...")
clear_cache()
inputs, outputs = prepare_model_for_inference(args, model, config, tokenizer, prompt_length, prompt)
all_token_ids = inputs["input_ids"].clone()
current_length = all_token_ids.shape[-1]
num_heads = config.num_key_value_heads
head_size = (
config.head_dim if hasattr(config, "head_dim") else config.hidden_size // config.num_attention_heads
)
has_eos = torch.zeros(batch_size, device=args.target_device, dtype=torch.bool)
# 0th entry will have prompt accelerator time, 1st entry onwards will have token generation accelerator time
accelerator_times = []
sampling_times = [] # cost to sample after each model run
wall_clock_start_time = time.perf_counter()
while current_length <= max_length:
# Run inference
accelerator_time_latency_s, outputs = run_inference(args, model, 1, inputs, outputs)
accelerator_times.append(accelerator_time_latency_s)
# Sample with argmax (greedy search)
sampling_start_time = time.perf_counter()
if outputs["logits"].shape[1] > 1:
prompt_end_indices = inputs["attention_mask"].sum(1) - 1
idxs = (
prompt_end_indices.unsqueeze(dim=1)
.repeat(1, config.vocab_size)
.view(batch_size, 1, config.vocab_size)
)
next_token_logits = torch.gather(outputs["logits"], 1, idxs).squeeze()
else:
next_token_logits = outputs["logits"][:, -1, :]
next_tokens = torch.argmax(next_token_logits, dim=-1)
# Check if we previously reached EOS token id or if generated token id is EOS token id
has_eos = has_eos | next_tokens == tokenizer.eos_token_id
# Determine which new tokens to add to list of all token ids
# Add EOS token ids for batch entries that ended early (ragged batching scenario where some batch entries ended early and some haven't)
tokens_to_add = next_tokens.masked_fill(has_eos, tokenizer.eos_token_id).reshape([batch_size, 1])
sampling_end_time = time.perf_counter()
sampling_times.append(sampling_end_time - sampling_start_time)
all_token_ids = torch.cat([all_token_ids, tokens_to_add], dim=-1)
# Return early if all batch entries have reached EOS token id
current_length += 1
if torch.all(has_eos) or current_length > max_length:
break
# Update inputs for next inference run
inputs["input_ids"] = tokens_to_add
inputs["attention_mask"] = torch.cat(
[inputs["attention_mask"], (~has_eos).to(torch.int64).reshape(batch_size, 1)], 1
)
inputs["position_ids"] = (
None
if "position_ids" not in inputs
else torch.max(inputs["position_ids"], dim=1)[0].reshape(batch_size, 1) + 1
)
# Set logits to zeros for next inference run and re-use memory buffer
if outputs["logits"].shape[1] != 1:
outputs["logits"] = outputs["logits"][:, :1, :].contiguous()
outputs["logits"].zero_()
# Update KV caches for next inference run
if args.engine == "pt":
# Update KV caches for PyTorch
inputs["past_key_values"] = outputs["past_key_values"]
elif not args.use_buffer_share:
# Update KV caches for ONNX Runtime if buffer sharing is not used
for i in range(config.num_hidden_layers):
inputs[f"past_key_values.{i}.key"] = outputs[f"present.{i}.key"]
inputs[f"past_key_values.{i}.value"] = outputs[f"present.{i}.value"]
new_sequence_length = inputs["attention_mask"].shape[1]
for i in range(config.num_hidden_layers):
present_key = torch.zeros(
batch_size,
num_heads,
new_sequence_length,
head_size,
device=args.target_device,
dtype=args.torch_dtype,
)
present_value = torch.zeros(
batch_size,
num_heads,
new_sequence_length,
head_size,
device=args.target_device,
dtype=args.torch_dtype,
)
outputs.update(
{
f"present.{i}.key": present_key.contiguous(),
f"present.{i}.value": present_value.contiguous(),
}
)
wall_clock_end_time = time.perf_counter()
# Filter out any anomaly accelerator times (e.g. for `torch.compile`)
accelerator_times.pop(0) # Remove prompt processing time
if args.anomaly_filtering:
anomaly_threshold_factor = 10
min_time_s = min(accelerator_times)
orig_size = len(accelerator_times)
accelerator_times = list(
filter(lambda acc_time: acc_time < anomaly_threshold_factor * min_time_s, accelerator_times)
)
new_size = len(accelerator_times)
logger.info(
f"Filtered out {orig_size - new_size} anomaly accelerator times that are {anomaly_threshold_factor}x greater than {min_time_s * 1000} ms..."
)
#######################################################
# Calculate sampling and first token generated metrics
#######################################################
# Calculate sampling metrics
avg_sampling_latency_s = sum(sampling_times) / len(sampling_times)
avg_sampling_latency_ms = avg_sampling_latency_s * 1000
avg_sampling_thrpt = batch_size * (1 / avg_sampling_latency_s)
logger.info(f"Average Latency of Sampling: {avg_sampling_latency_ms} ms")
logger.info(f"Average Throughput of Sampling: {avg_sampling_thrpt} tps")
# Calculate first token generated metrics
first_token_latency_s = accelerator_times[0]
first_token_latency_ms = first_token_latency_s * 1000
first_token_thrpt = batch_size * (1 / first_token_latency_s)
logger.info(f"Latency of First Token Generated: {first_token_latency_ms} ms")
logger.info(f"Throughput of First Token Generated: {first_token_thrpt} tps")
####################################################
# Calculate first `halfway` token generated metrics
####################################################
halfway = args.generation_length // 2
halfway_token_latency_s = sum(accelerator_times[:halfway]) / len(accelerator_times[:halfway])
halfway_token_latency_ms = halfway_token_latency_s * 1000
halfway_token_thrpt = batch_size * (1 / halfway_token_latency_s)
logger.info(f"Average Latency of First {halfway} Tokens Generated: {halfway_token_latency_ms} ms")
logger.info(f"Average Throughput of First {halfway} Tokens Generated: {halfway_token_thrpt} tps")
#########################################
# Calculate all tokens generated metrics
#########################################
all_token_latency_s = sum(accelerator_times) / len(accelerator_times)
all_token_latency_ms = all_token_latency_s * 1000
all_token_thrpt = batch_size * (1 / all_token_latency_s)
logger.info(
f"Average Latency of First {args.generation_length} Tokens Generated: {all_token_latency_ms} ms"
)
logger.info(f"Average Throughput of First {args.generation_length} Tokens Generated: {all_token_thrpt} tps")
###############################
# Calculate wall clock metrics
###############################
wall_clock_latency_s = wall_clock_end_time - wall_clock_start_time
wall_clock_thrpt = batch_size * ((prompt_length + args.generation_length) / wall_clock_latency_s)
logger.info(f"Wall-Clock Latency: {wall_clock_latency_s} s")
logger.info(
f"Wall-Clock Throughput: {batch_size * ((prompt_length + args.generation_length) / wall_clock_latency_s)} tps"
)
# Add metrics to CSV
logger.info("Adding results to CSV")
csv_metrics.extend(
[
avg_sampling_latency_ms,
avg_sampling_thrpt,
first_token_latency_ms,
first_token_thrpt,
halfway_token_latency_ms,
halfway_token_thrpt,
all_token_latency_ms,
all_token_thrpt,
wall_clock_latency_s,
wall_clock_thrpt,
]
)
all_csv_metrics.append(csv_metrics)
except: # noqa: E722
logger.info(f"Could not benchmark at batch size = {batch_size}, prompt length = {prompt_length}")
filename = f"benchmark_{args.engine}_e2e_{datetime.datetime.now():%Y-%m-%d_%H:%M:%S}.csv"
save_results(all_csv_metrics, filename, args.generation_length)
if __name__ == "__main__":
main()

5
onnxruntime/python/tools/transformers/models/llama/convert_to_onnx.py
View file

@ -1,3 +1,8 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
from __future__ import annotations
import argparse

5
onnxruntime/python/tools/transformers/models/llama/dist_settings.py
View file

@ -1,3 +1,8 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
import os
import torch.distributed as dist

194
onnxruntime/python/tools/transformers/models/llama/llama_inputs.py
View file

@ -1,8 +1,13 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
from __future__ import annotations
import numpy as np
import torch
from transformers import AutoConfig
from transformers import AutoConfig, AutoTokenizer
from onnxruntime import InferenceSession, OrtValue
@ -269,6 +274,8 @@ def convert_inputs_for_ort(
return ort_inputs
# Re-allocate KV caches from (batch_size, num_heads, past_sequence_length, head_size) to
# (batch_size, num_heads, max_sequence_length, head_size) for past-present buffer sharing
def enable_past_present_share_buffer(ort_inputs: dict, past_seq_len: int, max_seq_len: int):
for k, v in ort_inputs.items():
# Allocate new buffers with max_sequence_length for GQA
@ -281,8 +288,29 @@ def enable_past_present_share_buffer(ort_inputs: dict, past_seq_len: int, max_se
return ort_inputs
# Add IO bindings for execution providers
def add_io_bindings(
# Verify ONNX Runtime inputs with model
def verify_ort_inputs(model: InferenceSession, ort_inputs: dict):
# Check that all model inputs will be provided
model_inputs = set(map(lambda model_input: model_input.name, model.get_inputs()))
user_inputs = set(ort_inputs.keys())
missing_inputs = model_inputs - user_inputs
if len(missing_inputs):
print(f"The following model inputs are missing: {missing_inputs}")
raise Exception("There are missing inputs to the model. Please add them and try again.")
# Remove unnecessary inputs from model inputs
unnecessary_inputs = user_inputs - model_inputs
if len(unnecessary_inputs):
for unnecessary_input in unnecessary_inputs:
print(f"Removing unnecessary input '{unnecessary_input}' from user provided inputs")
del ort_inputs[unnecessary_input]
return ort_inputs
# Add IO bindings for execution providers using OrtValue
# Use when you need to run inference once or twice to save memory
def add_io_bindings_as_ortvalues(
model: InferenceSession, ort_inputs: dict, device: str, device_id: int, use_gqa: bool, kv_cache_ortvalues: dict
):
io_binding = model.io_binding()
@ -318,3 +346,163 @@ def add_io_bindings(
io_binding.bind_output(name, device_type=device, device_id=device_id)
return io_binding, kv_cache_ortvalues
# Add IO bindings for execution providers using PyTorch tensors
# Use when you need to run inference many times
def add_io_bindings_as_tensors(
model: InferenceSession, inputs: dict, outputs: dict, use_fp16: bool, use_buffer_share: bool
):
# Verify model inputs
inputs = verify_ort_inputs(model, inputs)
device = None
pt_to_np = {
"torch.int32": np.int32,
"torch.int64": np.int64,
"torch.float16": np.float16,
"torch.float32": np.float32,
}
# Bind inputs/outputs to IO binding
io_binding = model.io_binding()
for k, v in inputs.items():
io_binding.bind_input(
name=k,
device_type=v.device.type,
device_id=0 if v.device.type == "cpu" else v.device.index,
element_type=pt_to_np[repr(v.dtype)],
shape=tuple(v.shape),
buffer_ptr=v.data_ptr(),
)
device = v.device
for output in model.get_outputs():
name = output.name
if use_buffer_share and "present" in name:
# Bind KV cache outputs to KV cache inputs
v = inputs[name.replace("present", "past_key_values")]
io_binding.bind_output(
name=name,
device_type=v.device.type,
device_id=v.device.index,
element_type=np.float16,
shape=tuple(v.shape),
buffer_ptr=v.data_ptr(),
)
else:
v = outputs[name]
io_binding.bind_output(
name=name,
device_type=device.type,
device_id=0 if device.type == "cpu" else device.index,
element_type=(np.float16 if use_fp16 else np.float32),
shape=tuple(v.shape),
buffer_ptr=v.data_ptr(),
)
return io_binding
# Get actual inputs when using real data (instead of sample data) and initialize outputs
def get_initial_inputs_and_outputs(
config: AutoConfig,
tokenizer: AutoTokenizer,
requested_length: int,
prompt: list[str],
device: torch.device,
use_fp16: bool,
use_buffer_share: bool,
engine: str,
):
tokenizer.pad_token = "[PAD]"
encodings_dict = tokenizer.batch_encode_plus(prompt, padding=True)
torch_dtype = torch.float16 if use_fp16 else torch.float32
# input_ids: pad token id is 0
# attention_mask: pad token id is 0
# position_ids: pad token id is 1
input_ids = torch.tensor(encodings_dict["input_ids"], device=device, dtype=torch.int64)
attention_mask = torch.tensor(encodings_dict["attention_mask"], device=device, dtype=torch.int64)
position_ids = get_position_ids(attention_mask, use_past_kv=False)
# Check if tokenized prompt length matches the requested prompt length
tokenized_length = input_ids.shape[-1]
if tokenized_length > requested_length:
# Shorten the inputs from (batch_size, tokenized_length) to (batch_size, requested_length)
input_ids = input_ids[:, :requested_length]
attention_mask = attention_mask[:, :requested_length]
position_ids = get_position_ids(attention_mask, use_past_kv=False)
elif tokenized_length < requested_length:
# Lengthen the inputs from (batch_size, tokenized_length) to (batch_size, requested_length)
input_ids_first_col = input_ids[:, 0].unsqueeze(0).T
attention_mask_first_col = attention_mask[:, 0].unsqueeze(0).T
for _ in range(requested_length - tokenized_length):
input_ids = torch.hstack((input_ids_first_col, input_ids))
attention_mask = torch.hstack((attention_mask_first_col, attention_mask))
position_ids = get_position_ids(attention_mask, use_past_kv=False)
tokenized_length = input_ids.shape[-1]
assert tokenized_length == requested_length
# Create inputs
inputs = {
"input_ids": input_ids.contiguous() if engine == "ort" else input_ids,
"attention_mask": attention_mask.contiguous() if engine == "ort" else attention_mask,
"position_ids": position_ids.contiguous() if engine == "ort" else position_ids,
}
if engine != "ort":
inputs["past_key_values"] = []
# Get shape of KV cache inputs
batch_size, sequence_length = input_ids.shape
max_sequence_length = config.max_position_embeddings
num_heads = config.num_key_value_heads
head_size = config.head_dim if hasattr(config, "head_dim") else config.hidden_size // config.num_attention_heads
# Create KV cache inputs
for i in range(config.num_hidden_layers):
past_key = torch.zeros(
batch_size,
num_heads,
max_sequence_length if use_buffer_share else 0,
head_size,
device=device,
dtype=torch_dtype,
)
past_value = torch.zeros(
batch_size,
num_heads,
max_sequence_length if use_buffer_share else 0,
head_size,
device=device,
dtype=torch_dtype,
)
if engine == "ort":
inputs.update(
{
f"past_key_values.{i}.key": past_key.contiguous(),
f"past_key_values.{i}.value": past_value.contiguous(),
}
)
else:
inputs["past_key_values"].append((past_key, past_value))
outputs = None
if engine == "ort":
# Create outputs
logits = torch.zeros(batch_size, sequence_length, config.vocab_size, device=device, dtype=torch_dtype)
outputs = {"logits": logits.contiguous()}
if not use_buffer_share:
for i in range(config.num_hidden_layers):
present_key = torch.zeros(
batch_size, num_heads, sequence_length, head_size, device=device, dtype=torch_dtype
)
present_value = torch.zeros(
batch_size, num_heads, sequence_length, head_size, device=device, dtype=torch_dtype
)
outputs.update(
{f"present.{i}.key": present_key.contiguous(), f"present.{i}.value": present_value.contiguous()}
)
return inputs, outputs

9
onnxruntime/python/tools/transformers/models/llama/llama_parity.py
View file

@ -1,3 +1,8 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
from __future__ import annotations
import argparse
@ -10,7 +15,7 @@ import torch
from benchmark_helper import setup_logger
from dist_settings import get_rank, get_size
from llama_inputs import (
add_io_bindings,
add_io_bindings_as_ortvalues,
convert_inputs_for_ort,
get_merged_sample_with_past_kv_inputs,
get_sample_inputs,
@ -123,7 +128,7 @@ def verify_parity(
# Add IO bindings for non-CPU execution providers
if args.execution_provider != "cpu":
io_binding, kv_cache_ortvalues = add_io_bindings(
io_binding, kv_cache_ortvalues = add_io_bindings_as_ortvalues(
ort_model,
inputs,
args.execution_provider,

5
onnxruntime/python/tools/transformers/models/llama/llama_torch.py
View file

@ -1,3 +1,8 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
import logging
import os

11
onnxruntime/python/tools/transformers/models/llama/prompts.json Normal file
View file

File diff suppressed because one or more lines are too long

5
onnxruntime/python/tools/transformers/models/llama/quant_kv_dataloader.py
View file

@ -1,3 +1,8 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
import argparse
import numpy as np