saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-10 17:37:14 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

GPT2 with one step beam search (#7163)

Browse source 
* beam search refactoring checkin
* add factory class and deduplicate code
* one step beam search works on gpu

Co-authored-by: Xiaoyu Liu <xiaoyu@xiaoyu-VM.z4vh1dzj5eoevgybsksdpz2izh.jx.internal.cloudapp.net>
This commit is contained in:
Xiaoyu Liu 2021-04-20 06:23:52 -07:00 committed by GitHub
parent 1a3ddf0714
commit 913ea8264b
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 1291 additions and 8 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

41
onnxruntime/python/tools/transformers/convert_to_onnx.py
View file

@ -25,8 +25,9 @@ import json
from pathlib import Path
from packaging import version
from transformers import AutoConfig
from gpt2_helper import Gpt2Helper, MODEL_CLASSES, DEFAULT_TOLERANCE, PRETRAINED_GPT2_MODELS
from gpt2_tester import Gpt2Tester
from gpt2_helper import DEFAULT_TOLERANCE, PRETRAINED_GPT2_MODELS
from gpt2_beamsearch_helper import Gpt2HelperFactory, MODEL_CLASSES
from gpt2_beamsearch_tester import Gpt2TesterFactory
from quantize_helper import QuantizeHelper
from benchmark_helper import create_onnxruntime_session, setup_logger, prepare_environment, Precision
@ -99,6 +100,9 @@ def parse_arguments():
parser.add_argument('-e', '--use_external_data_format', required=False, action='store_true')
parser.set_defaults(use_external_data_format=False)
parser.add_argument('--batch_size', required=False, type=int, default=1, help='Batch size for GPT model with beam search')
parser.add_argument('--beam_size', required=False, type=int, default=4, help='Beam size for beam search')
args = parser.parse_args()
return args
@ -134,8 +138,14 @@ def main():
assert not args.output.endswith('.onnx'), "output shall be a directory for --use_external_data_format"
model_class = MODEL_CLASSES[args.model_class][0]
model_type = "beam_search_step" if args.model_class == "GPT2LMHeadModel_BeamSearchStep" else "default"
gpt2helper = Gpt2HelperFactory.create_helper(model_type)
gpt2tester = Gpt2TesterFactory.create_tester(model_type)
config = AutoConfig.from_pretrained(args.model_name_or_path, cache_dir=cache_dir)
model = model_class.from_pretrained(args.model_name_or_path, config=config, cache_dir=cache_dir)
if model_type == 'beam_search_step':
model = model_class.from_pretrained(args.model_name_or_path, config=config, batch_size=args.batch_size, beam_size=args.beam_size, cache_dir=cache_dir)
else:
model = model_class.from_pretrained(args.model_name_or_path, config=config, cache_dir=cache_dir)
device = torch.device("cuda:0" if args.use_gpu else "cpu")
model.eval().to(device)
@ -143,7 +153,7 @@ def main():
if (not args.use_external_data_format) and (config.n_layer > 24):
logger.info(f"Try --use_external_data_format when model size > 2GB")
onnx_model_paths = Gpt2Helper.get_onnx_paths(output_dir,
onnx_model_paths = gpt2helper.get_onnx_paths(output_dir,
args.model_name_or_path,
args.model_class,
new_folder=args.use_external_data_format)
@ -152,7 +162,7 @@ def main():
logger.info(f"Exporting ONNX model to {raw_onnx_model}")
use_padding = MODEL_CLASSES[args.model_class][2]
Gpt2Helper.export_onnx(model,
gpt2helper.export_onnx(model,
device,
raw_onnx_model,
args.verbose,
@ -164,7 +174,7 @@ def main():
output_path = onnx_model_paths[str(args.precision) if args.precision != Precision.INT8 else 'fp32']
logger.info(f"Optimizing model to {output_path}")
Gpt2Helper.optimize_onnx(raw_onnx_model, output_path, args.precision == Precision.FLOAT16,
gpt2helper.optimize_onnx(raw_onnx_model, output_path, args.precision == Precision.FLOAT16,
model.config.num_attention_heads, model.config.hidden_size,
args.use_external_data_format)
else:
@ -186,7 +196,7 @@ def main():
session = create_onnxruntime_session(output_path, args.use_gpu, enable_all_optimization=True, verbose=args.verbose)
if session is not None:
Gpt2Helper.test_parity(session,
gpt2helper.test_parity(session,
model,
device,
args.precision == Precision.FLOAT16,
@ -229,9 +239,24 @@ def main():
else:
inputs = {"input_ids": input_ids}
if model_type == "beam_search_step":
beam_select_idx = torch.zeros([1, input_ids.shape[0]]).long()
input_log_probs = torch.zeros([input_ids.shape[0], 1])
input_unfinished_sents = torch.ones(
[input_ids.shape[0], 1], dtype=torch.bool
)
inputs.update(
{
"beam_select_idx": beam_select_idx,
"input_log_probs": input_log_probs,
"input_unfinished_sents": input_unfinished_sents,
}
)
test_inputs.append(inputs)
Gpt2Tester.test_generation(session,
gpt2tester.test_generation(session,
model,
device,
test_inputs,

824
onnxruntime/python/tools/transformers/gpt2_beamsearch_helper.py Normal file
View file

@ -0,0 +1,824 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
# This script helps onnx conversion and validation for GPT2 model with past state.
import os
import logging
import torch
import onnx
import random
import numpy
import time
import re
from pathlib import Path
from typing import List, Dict, Tuple, Union
from transformers import GPT2LMHeadModel, GPT2Config
from benchmark_helper import Precision
from gpt2_helper import Gpt2Helper, Gpt2Inputs, GPT2ModelNoPastState, MyGPT2Model, MyGPT2LMHeadModel, MyGPT2LMHeadModel_NoPadding
logger = logging.getLogger(__name__)
class Gpt2HelperFactory:
@staticmethod
def create_helper(helper_type="default"):
helpers = {
"default": Gpt2Helper,
"beam_search_step": Gpt2BeamSearchHelper,
}
w = helpers[helper_type]
return w
class GPT2LMHeadModel_BeamSearchStep(GPT2LMHeadModel):
"""Here we wrap a class for Onnx model conversion for GPT2LMHeadModel with past state and one
step beam search."""
def __init__(self, config, batch_size, beam_size):
super().__init__(config)
self.config.batch_size = batch_size
self.config.beam_size = beam_size
def forward(
self,
input_ids,
position_ids,
attention_mask,
beam_select_idx,
input_log_probs,
input_unfinished_sents,
prev_step_results,
prev_step_scores,
*past,
):
input_ids = input_ids.view(self.config.batch_size, -1, input_ids.size(-1))
past = [past[i].index_select(1, beam_select_idx[0]) for i in range(len(past))]
result = super().forward(
input_ids.view(-1, input_ids.size(-1)),
position_ids=position_ids,
attention_mask=attention_mask,
past_key_values=past,
return_dict=False,
)
logits_flat, present_flat = MyGPT2Model.post_process(result, self.config.n_layer)
next_token_logits = logits_flat[:, -1].view(
self.config.batch_size, -1, logits_flat.size(-1)
)
next_token_log_probs = torch.log_softmax(next_token_logits, dim=-1)
next_token_log_probs, next_token_ids = torch.topk(
next_token_log_probs, self.config.beam_size, dim=-1, largest=True, sorted=True
)
# finished sentences is always with EOS, and all but the first one has -inf, so that they will be automatically dropped in the round of beam search.
finished_sents = ~input_unfinished_sents
next_token_log_probs.masked_fill_(finished_sents.unsqueeze(-1), -numpy.inf)
next_token_log_probs[..., 0].masked_fill_(finished_sents, 0)
next_token_ids.masked_fill_(
finished_sents.unsqueeze(-1), self.config.eos_token_id
)
output_log_probs = input_log_probs.unsqueeze(-1) + next_token_log_probs
# select N sequences from beams of each input, sorted by sequence probability
output_log_probs = output_log_probs.view(
self.config.batch_size, -1
) # shape=(batch, beam_size^2)
output_log_probs, selected_index_flat = output_log_probs.topk(
self.config.beam_size, dim=-1, largest=True, sorted=True
) # output shape=(batch, beam_size)
# select the correspondent sentences/next tokens
selected_input_seq = selected_index_flat // self.config.beam_size
next_token_ids = next_token_ids.view(self.config.batch_size, -1).gather(
-1, selected_index_flat
)
prev_step_results = prev_step_results.view(
self.config.batch_size, -1, prev_step_results.size(-1)
)
prev_step_results = prev_step_results.gather(
1, selected_input_seq.unsqueeze(-1).repeat(1, 1, prev_step_results.size(-1))
)
output_unfinished_sents = input_unfinished_sents.gather(1, selected_input_seq)
output_unfinished_sents = (
output_unfinished_sents
& next_token_ids.ne(self.config.eos_token_id)
)
# get the next full input_ids
current_step_results = torch.cat(
[prev_step_results, next_token_ids.unsqueeze(-1)], dim=-1
).contiguous()
prev_step_scores = prev_step_scores.view(
self.config.batch_size, -1, prev_step_scores.size(-1)
)
prev_step_scores = prev_step_scores.gather(
1, selected_input_seq.unsqueeze(-1).repeat(1, 1, prev_step_scores.size(-1))
)
current_step_scores = torch.cat(
[prev_step_scores, output_log_probs.unsqueeze(-1)], dim=-1
).contiguous()
return (
next_token_ids,
present_flat,
selected_input_seq,
output_log_probs,
output_unfinished_sents,
current_step_results.view(self.config.batch_size * self.config.beam_size, -1),
current_step_scores.view(self.config.batch_size * self.config.beam_size, -1),
)
# Maps model class name to a tuple of model class, name of first output and use padding or not
MODEL_CLASSES = {
'GPT2LMHeadModel': (MyGPT2LMHeadModel, 'logits', True),
'GPT2LMHeadModel_NoPadding': (MyGPT2LMHeadModel_NoPadding, 'logits', False),
'GPT2Model': (MyGPT2Model, 'last_state', True),
"GPT2LMHeadModel_BeamSearchStep": (GPT2LMHeadModel_BeamSearchStep, "last_state", True), # defined in gpt2_beamsearch_helper.py
}
class Gpt2BeamSearchInputs(Gpt2Inputs):
def __init__(
self,
input_ids,
position_ids,
attention_mask,
past,
beam_select_idx=None,
input_log_probs=None,
input_unfinished_sents=None,
prev_step_results=None,
prev_step_scores=None,
):
super().__init__(input_ids, position_ids, attention_mask, past)
self.prev_step_results: torch.LongTensor = prev_step_results
self.prev_step_scores: Union[torch.FloatTensor, torch.HalfTensor, torch.cuda.FloatTensor] = prev_step_scores
if beam_select_idx is None:
self.beam_select_idx: torch.LongTensor = torch.zeros(
[1, len(input_ids)]
).long()
else:
self.beam_select_idx: torch.LongTensor = beam_select_idx
self.input_log_probs: Union[torch.FloatTensor, torch.HalfTensor, torch.cuda.FloatTensor] = input_log_probs
self.input_unfinished_sents: torch.ByteTensor = input_unfinished_sents
def to_list(self) -> List:
input_list = [
v
for v in [
self.input_ids,
self.position_ids,
self.attention_mask,
self.beam_select_idx,
self.input_log_probs,
self.input_unfinished_sents,
self.prev_step_results,
self.prev_step_scores
]
if v is not None
]
if self.past:
input_list.extend(self.past)
return input_list
def to_fp32(self):
gpt2_inputs = super().to_fp32()
return Gpt2BeamSearchInputs(
gpt2_inputs.input_ids,
gpt2_inputs.position_ids,
gpt2_inputs.attention_mask,
gpt2_inputs.past,
self.beam_select_idx,
self.input_log_probs.to(dtype=torch.float32),
self.input_unfinished_sents,
self.prev_step_results,
self.prev_step_scores.to(dtype=torch.float32),
)
class Gpt2BeamSearchHelper(Gpt2Helper):
"""A helper class for Gpt2 model conversion, inference and verification."""
@staticmethod
def get_dummy_inputs(batch_size: int,
past_sequence_length: int,
sequence_length: int,
num_attention_heads: int,
hidden_size: int,
num_layer: int,
vocab_size: int,
device: torch.device,
float16: bool = False,
has_position_ids: bool = True,
has_attention_mask: bool = True) -> Gpt2BeamSearchInputs:
"""Create random inputs for GPT2 model.
Returns torch tensors of input_ids, position_ids, attention_mask and a list of past state tensors.
"""
gpt2_dummy_inputs = Gpt2Helper.get_dummy_inputs(
batch_size,
past_sequence_length,
sequence_length,
num_attention_heads,
hidden_size,
num_layer,
vocab_size,
device,
float16,
has_position_ids,
has_attention_mask
)
float_type = torch.float16 if float16 else torch.float32
beam_select_idx = torch.zeros([1, batch_size], device=device).long()
input_log_probs = torch.zeros([batch_size, 1], dtype=float_type, device=device)
input_unfinished_sents = torch.ones(
[batch_size, 1], dtype=torch.bool, device=device
)
prev_step_results = torch.randint(
low=0,
high=vocab_size - 1,
size=(batch_size, sequence_length),
dtype=torch.int64,
device=device,
)
prev_step_scores = torch.zeros([batch_size, 1], dtype=float_type, device=device)
return Gpt2BeamSearchInputs(
gpt2_dummy_inputs.input_ids,
gpt2_dummy_inputs.position_ids,
gpt2_dummy_inputs.attention_mask,
gpt2_dummy_inputs.past,
beam_select_idx,
input_log_probs,
input_unfinished_sents,
prev_step_results,
prev_step_scores,
)
@staticmethod
def get_output_shapes(batch_size: int,
context_len: int,
past_sequence_length: int,
sequence_length: int,
beam_size: int,
step: int,
config: GPT2Config,
model_class: str = "GPT2LMHeadModel") -> Dict[str, List[int]]:
"""Returns a dictionary with output name as key, and shape as value."""
num_attention_heads = config.num_attention_heads
hidden_size = config.hidden_size
num_layer = config.num_hidden_layers
vocab_size = config.vocab_size
output_name = MODEL_CLASSES[model_class][1]
last_state_shape = [batch_size, beam_size]
if step == 0:
present_state_shape = [
2,
batch_size,
num_attention_heads,
past_sequence_length + sequence_length,
int(hidden_size / num_attention_heads),
]
else:
present_state_shape = [
2,
batch_size * beam_size,
num_attention_heads,
past_sequence_length + sequence_length,
int(hidden_size / num_attention_heads),
]
output_shapes = {output_name: last_state_shape}
for i in range(num_layer):
output_shapes["present_" + str(i)] = present_state_shape
output_shapes["output_selected_indices"] = [1, batch_size * beam_size]
output_shapes["output_log_probs"] = [batch_size, beam_size]
output_shapes["output_unfinished_sents"] = [batch_size, beam_size]
output_shapes["current_step_results"] = [batch_size * beam_size, past_sequence_length + sequence_length + 1]
output_shapes["current_step_scores"] = [batch_size * beam_size, past_sequence_length + sequence_length - context_len + 2]
return output_shapes
@staticmethod
def get_output_buffers(
output_shapes, device, is_float16=False
):
"""Returns a dictionary of output name as key, and 1D tensor as value. The tensor has enough space for given shape."""
data_type = torch.float16 if is_float16 else torch.float32
output_buffers = {}
for name, shape in output_shapes.items():
if (
name == "output_selected_indices"
or name == "current_step_results"
or name == "last_state"
):
output_buffers[name] = torch.empty(
numpy.prod(shape), dtype=torch.long, device=device
)
elif name == "output_unfinished_sents":
output_buffers[name] = torch.empty(
numpy.prod(shape), dtype=torch.bool, device=device
)
else:
output_buffers[name] = torch.empty(
numpy.prod(shape), dtype=data_type, device=device
)
return output_buffers
@staticmethod
def compare_outputs(torch_outputs, ort_outputs, rtol=1e-03, atol=1e-03):
"""Returns True if torch and ORT outputs are close for given thresholds, and False otherwise."""
is_close = numpy.allclose(
ort_outputs[-4], torch_outputs[-4].cpu().numpy(), rtol=rtol, atol=atol
)
logger.debug(
f"PyTorch and OnnxRuntime output 0 (last_state) are close: {is_close}"
)
is_all_close = is_close
num_layers = len(ort_outputs) - 6
for layer in range(num_layers):
is_close = numpy.allclose(
ort_outputs[1 + layer],
torch_outputs[1][layer].cpu().numpy(),
rtol=rtol,
atol=atol,
)
logger.debug(
f"PyTorch and OnnxRuntime layer {layer} state (present_{layer}) are close:{is_close}"
)
is_all_close = is_all_close and is_close
if not is_all_close:
max_abs_diff = Gpt2BeamSearchHelper.diff_outputs(torch_outputs, ort_outputs)
logger.info(
f"PyTorch and OnnxRuntime results are not all close: max_abs_diff={max_abs_diff:.5f}"
)
return is_all_close
@staticmethod
def export_onnx(model,
device,
onnx_model_path: str,
verbose: bool = False,
use_external_data_format: bool = False,
has_position_ids: bool = True,
has_attention_mask: bool = True):
"""Export GPT-2 model with past state to ONNX model."""
config: GPT2Config = model.config
num_layer = config.n_layer
dummy_inputs = Gpt2BeamSearchHelper.get_dummy_inputs(batch_size=1,
past_sequence_length=1,
sequence_length=1,
num_attention_heads=config.num_attention_heads,
hidden_size=config.hidden_size,
num_layer=num_layer,
vocab_size=config.vocab_size,
device=device,
float16=False,
has_position_ids=has_position_ids,
has_attention_mask=has_attention_mask)
input_list = dummy_inputs.to_list()
with torch.no_grad():
# outputs = model(input_ids, position_id, attention_mask, beam_select_idx, past)
outputs = model(*input_list)
past_names = [f"past_{i}" for i in range(num_layer)]
present_names = [f"present_{i}" for i in range(num_layer)]
output_names = ["last_state"] + present_names
output_names += [
"output_selected_indices",
"output_log_probs",
"output_unfinished_sents",
"current_step_results",
"current_step_scores",
]
# Shape of input tensors:
# input_ids: (batch_size, seq_len)
# past_{i}: (2, batch_size, num_heads, past_seq_len, hidden_size/num_heads)
# attention_mask: (batch_size, past_seq_len + seq_len)
# Shape of output tensors:
# last_state: (batch_size, seq_len, hidden_size)
# or logits: (batch_size, seq_len, vocab_size)
# present_{i}: (2, batch_size, num_heads, past_seq_len + seq_len, hidden_size/num_heads)
dynamic_axes = {
"input_ids": {0: "batch_size", 1: "seq_len"},
output_names[0]: {0: "batch_size", 1: "seq_len"},
}
for name in past_names:
dynamic_axes[name] = {1: "batch_size", 3: "past_seq_len"}
for name in present_names:
dynamic_axes[name] = {1: "batch_size", 3: "total_seq_len"}
input_names = ["input_ids"]
dynamic_axes["position_ids"] = {0: "batch_size", 1: "seq_len"}
input_names.append("position_ids")
dynamic_axes["attention_mask"] = {0: "batch_size", 1: "total_seq_len"}
input_names.append("attention_mask")
dynamic_axes["beam_select_idx"] = {1: "batch_size"}
input_names.append("beam_select_idx")
dynamic_axes["input_log_probs"] = {0: "batch_size", 1: "beam_size"}
input_names.append("input_log_probs")
dynamic_axes["input_unfinished_sents"] = {0: "batch_size", 1: "beam_size"}
input_names.append("input_unfinished_sents")
dynamic_axes["prev_step_results"] = {0: "batch_size", 1: "total_seq_len"}
input_names.append("prev_step_results")
dynamic_axes["prev_step_scores"] = {0: "batch_size", 1: "total_seq_len"}
input_names.append("prev_step_scores")
input_names.extend(past_names)
logger.info(
f"Shapes: input_ids={dummy_inputs.input_ids.shape} past={dummy_inputs.past[0].shape} output={outputs[0].shape} present={outputs[1][0].shape}"
)
Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
torch.onnx.export(
model,
args=tuple(input_list),
f=onnx_model_path,
input_names=input_names,
output_names=output_names,
example_outputs=outputs,
dynamic_axes=dynamic_axes,
opset_version=12,
do_constant_folding=True,
use_external_data_format=use_external_data_format,
verbose=verbose,
)
@staticmethod
def onnxruntime_inference(ort_session, inputs: Gpt2BeamSearchInputs, total_runs: int = 0):
"""Run inference of ONNX model, and returns average latency in ms when total_runs > 0 besides outputs."""
logger.debug(f"start onnxruntime_inference")
ort_inputs = {
"input_ids": numpy.ascontiguousarray(inputs.input_ids.cpu().numpy())
}
if inputs.position_ids is not None:
ort_inputs["position_ids"] = numpy.ascontiguousarray(
inputs.position_ids.cpu().numpy()
)
if inputs.attention_mask is not None:
ort_inputs["attention_mask"] = numpy.ascontiguousarray(
inputs.attention_mask.cpu().numpy()
)
if inputs.beam_select_idx is not None:
ort_inputs["beam_select_idx"] = numpy.ascontiguousarray(
inputs.beam_select_idx.cpu().numpy()
)
if inputs.input_log_probs is not None:
ort_inputs["input_log_probs"] = numpy.ascontiguousarray(
inputs.input_log_probs.cpu().numpy()
)
if inputs.input_unfinished_sents is not None:
ort_inputs["input_unfinished_sents"] = numpy.ascontiguousarray(
inputs.input_unfinished_sents.cpu().numpy()
)
if inputs.prev_step_results is not None:
ort_inputs["prev_step_results"] = numpy.ascontiguousarray(
inputs.prev_step_results.cpu().numpy()
)
if inputs.prev_step_scores is not None:
ort_inputs["prev_step_scores"] = numpy.ascontiguousarray(
inputs.prev_step_scores.cpu().numpy()
)
if inputs.past is not None:
for i, past_i in enumerate(inputs.past):
ort_inputs[f"past_{i}"] = numpy.ascontiguousarray(past_i.cpu().numpy())
ort_outputs = ort_session.run(None, ort_inputs)
if total_runs == 0:
return ort_outputs
latency = []
for _ in range(total_runs):
start = time.time()
ort_outputs = ort_session.run(None, ort_inputs)
latency.append(time.time() - start)
average_latency = sum(latency) * 1000 / len(latency)
logger.debug(
"OnnxRuntime Inference time = {} ms".format(format(average_latency, ".2f"))
)
return ort_outputs, average_latency
@staticmethod
def prepare_io_binding(ort_session,
input_ids,
position_ids,
attention_mask,
past,
output_buffers,
output_shapes,
beam_select_idx=None,
input_log_probs=None,
input_unfinished_sents=None,
prev_step_results=None,
prev_step_scores=None):
"""Returnas IO binding object for a session."""
# Bind inputs and outputs to onnxruntime session
io_binding = Gpt2Helper.prepare_io_binding(ort_session, input_ids, position_ids, attention_mask, past, output_buffers, output_shapes)
# Bind inputs
data_type = output_buffers[ort_session.get_outputs()[1].name].dtype
float_type = numpy.float16 if data_type == torch.float16 else numpy.float32
if past is not None:
for i, past_i in enumerate(past):
assert past_i.is_contiguous()
data_ptr = past_i.data_ptr()
if data_ptr == 0:
# When past_sequence_length is 0, its data_ptr will be zero. IO Binding asserts that data_ptr shall not be zero.
# Here we workaround and pass data pointer of input_ids. Actual data is not used for past so it does not matter.
data_ptr = input_ids.data_ptr()
io_binding.bind_input(f'past_{i}', past_i.device.type, 0, float_type, list(past_i.size()), data_ptr)
if attention_mask is not None:
assert attention_mask.is_contiguous()
io_binding.bind_input('attention_mask', attention_mask.device.type, 0, float_type,
list(attention_mask.size()), attention_mask.data_ptr())
if beam_select_idx is not None:
assert beam_select_idx.is_contiguous()
io_binding.bind_input(
"beam_select_idx",
beam_select_idx.device.type,
0,
numpy.longlong,
list(beam_select_idx.size()),
beam_select_idx.data_ptr(),
)
if input_log_probs is not None:
assert input_log_probs.is_contiguous()
io_binding.bind_input(
"input_log_probs",
input_log_probs.device.type,
0,
float_type,
list(input_log_probs.size()),
input_log_probs.data_ptr(),
)
if input_unfinished_sents is not None:
assert input_unfinished_sents.is_contiguous()
io_binding.bind_input(
"input_unfinished_sents",
input_unfinished_sents.device.type,
0,
numpy.bool,
list(input_unfinished_sents.size()),
input_unfinished_sents.data_ptr(),
)
if prev_step_results is not None:
assert prev_step_results.is_contiguous()
io_binding.bind_input(
"prev_step_results",
prev_step_results.device.type,
0,
numpy.longlong,
list(prev_step_results.size()),
prev_step_results.data_ptr(),
)
if prev_step_scores is not None:
assert prev_step_scores.is_contiguous()
io_binding.bind_input(
"prev_step_scores",
prev_step_scores.device.type,
0,
float_type,
list(prev_step_scores.size()),
prev_step_scores.data_ptr(),
)
# Bind outputs
for output in ort_session.get_outputs():
output_name = output.name
output_buffer = output_buffers[output_name]
logger.debug(
f"{output_name} device type={output_buffer.device.type} shape={list(output_buffer.size())}"
)
if (
output_name == "output_selected_indices"
or output_name == "last_state"
or output_name == "current_step_results"
):
io_binding.bind_output(
output_name,
output_buffer.device.type,
0,
numpy.longlong,
output_shapes[output_name],
output_buffer.data_ptr(),
)
elif output_name == "output_unfinished_sents":
io_binding.bind_output(
output_name,
output_buffer.device.type,
0,
numpy.bool,
output_shapes[output_name],
output_buffer.data_ptr(),
)
else:
io_binding.bind_output(
output_name,
output_buffer.device.type,
0,
float_type,
output_shapes[output_name],
output_buffer.data_ptr(),
)
return io_binding
@staticmethod
def onnxruntime_inference_with_binded_io(ort_session,
inputs: Gpt2BeamSearchInputs,
output_buffers: Dict[str, torch.Tensor],
output_shapes: Dict[str, List[int]],
total_runs: int = 0,
return_numpy: bool = True,
include_copy_output_latency: bool = False):
"""Inference with IO binding. Returns outputs, and optional latency when total_runs > 0.
"""
logger.debug(f"start onnxruntime_inference_with_binded_io")
# Bind inputs and outputs to onnxruntime session
io_binding = Gpt2BeamSearchHelper.prepare_io_binding(
ort_session,
inputs.input_ids,
inputs.position_ids,
inputs.attention_mask,
inputs.past,
output_buffers,
output_shapes,
inputs.beam_select_idx,
inputs.input_log_probs,
inputs.input_unfinished_sents,
inputs.prev_step_results,
inputs.prev_step_scores,
)
# Run onnxruntime with io binding
ort_session.run_with_iobinding(io_binding)
# Copy results to cpu for verification
ort_outputs = Gpt2BeamSearchHelper.get_outputs_from_io_binding_buffer(
ort_session, output_buffers, output_shapes, return_numpy
)
if total_runs == 0:
return ort_outputs
latency = []
for _ in range(total_runs):
start = time.time()
# Run onnxruntime with io binding
ort_session.run_with_iobinding(io_binding)
if include_copy_output_latency:
_ = Gpt2BeamSearchHelper.get_outputs_from_io_binding_buffer(
ort_session, output_buffers, output_shapes, return_numpy
)
latency.append(time.time() - start)
average_latency = sum(latency) * 1000 / len(latency)
logger.debug(
"OnnxRuntime with IO binding inference time = {} ms".format(
format(average_latency, ".2f")
)
)
return ort_outputs, average_latency
@staticmethod
def test_parity(ort_session,
model,
device,
is_float16=False,
rtol=5e-4,
atol=5e-4,
total_test_cases=100,
use_io_binding=True,
model_class="GPT2LMHeadModel_BeamSearchStep",
has_position_ids=True,
has_attention_mask=True):
"""Generate random inputs and compare the results of PyTorch and Onnx Runtime."""
config: GPT2Config = model.config
logger.info(
f"Running parity test (rtol={rtol}, atol={atol}, test_cases={total_test_cases}, use_io_binding={use_io_binding} model_class={model_class} is_float16={is_float16}) ..."
)
max_batch_size = 1
max_past_seq_len = 4 # Do not use large number here for higher chance of hitting empty past (past_seq_len=0)
max_seq_len = 2
beam_size = 4
output_buffers = None
if use_io_binding:
max_output_shapes = Gpt2BeamSearchHelper.get_output_shapes(
max_batch_size,
max_past_seq_len,
max_past_seq_len,
max_seq_len,
beam_size,
0,
config,
model_class,
)
output_buffers = Gpt2BeamSearchHelper.get_output_buffers(
max_output_shapes, device, is_float16
)
passed_test_cases = 0
for _ in range(total_test_cases):
sequence_length = random.randint(1, max_seq_len)
past_sequence_length = random.randint(0, max_past_seq_len)
batch_size = random.randint(1, max_batch_size)
logger.debug(
f"Running parity test for batch_size={batch_size} past_sequence_length={past_sequence_length}..."
)
dummy_inputs = Gpt2BeamSearchHelper.get_dummy_inputs(
batch_size,
past_sequence_length,
sequence_length,
config.num_attention_heads,
config.hidden_size,
config.n_layer,
config.vocab_size,
device,
is_float16,
has_position_ids,
has_attention_mask
)
outputs = Gpt2BeamSearchHelper.pytorch_inference(model, dummy_inputs)
if use_io_binding:
ort_outputs = Gpt2BeamSearchHelper.onnxruntime_inference(
ort_session, dummy_inputs
)
else:
output_shapes = Gpt2BeamSearchHelper.get_output_shapes(
batch_size,
past_sequence_length,
past_sequence_length,
sequence_length,
beam_size,
0,
config,
model_class,
)
ort_outputs = Gpt2BeamSearchHelper.onnxruntime_inference_with_binded_io(
ort_session, dummy_inputs, output_buffers, output_shapes
)
is_all_close = Gpt2BeamSearchHelper.compare_outputs(
outputs, ort_outputs, rtol=rtol, atol=atol
)
if is_all_close:
passed_test_cases += 1
logger.info(f"Parity Test Cases={total_test_cases}; Passed={passed_test_cases}")
if passed_test_cases > 0.95 * total_test_cases:
logger.info(
f"Parity is good: passed rate={int(passed_test_cases*100/total_test_cases):.0f}%"
)
return passed_test_cases == total_test_cases
@staticmethod
def torchscript(model, config, device, has_position_ids=True, has_attention_mask=True):
"""JIT trace for TorchScript."""
input_list = Gpt2BeamSearchHelper.get_dummy_inputs(
batch_size=1,
past_sequence_length=1,
sequence_length=1,
num_attention_heads=config.num_attention_heads,
hidden_size=config.hidden_size,
num_layer=config.n_layer,
vocab_size=config.vocab_size,
device=device,
float16=False,
has_position_ids=has_position_ids,
has_attention_mask=has_attention_mask,
).to_list()
return torch.jit.trace(model, input_list)

434
onnxruntime/python/tools/transformers/gpt2_beamsearch_tester.py Normal file
View file

@ -0,0 +1,434 @@
# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License. See License.txt in the project root for
# license information.
# --------------------------------------------------------------------------
# This script helps evaluation of GPT-2 model.
import os
import logging
import torch
import random
import numpy
import time
import timeit
import math
import statistics
from pathlib import Path
from gpt2_tester import Gpt2Tester, Gpt2Metric
from gpt2_beamsearch_helper import Gpt2BeamSearchHelper, Gpt2BeamSearchInputs
from benchmark_helper import Precision
logger = logging.getLogger(__name__)
class Gpt2TesterFactory:
@staticmethod
def create_tester(tester_type="default"):
testers = {
"default": Gpt2Tester,
"beam_search_step": Gpt2BeamSearchTester,
}
w = testers[tester_type]
return w
class Gpt2BeamSearchTester(Gpt2Tester):
def __init__(self,
input_ids,
position_ids,
attention_mask,
beam_select_idx,
input_log_probs,
input_unfinished_sents,
prev_step_results,
prev_step_scores,
num_attention_heads,
hidden_size,
num_layer,
beam_size,
device,
is_fp16=False,
top_k=20,
top_k_required_order=False,
):
super().__init__(
input_ids,
position_ids,
attention_mask,
num_attention_heads,
hidden_size,
num_layer,
device,
is_fp16,
top_k,
top_k_required_order
)
self.input_length = input_ids.shape[-1]
self.n_layer = num_layer
self.beam_size = beam_size
self.beam_select_idx = beam_select_idx.to(device)
float_type = torch.float16 if is_fp16 else torch.float32
self.input_log_probs = input_log_probs.type(float_type).to(device)
self.input_unfinished_sents = input_unfinished_sents.to(device)
self.prev_step_results = prev_step_results.to(device)
self.prev_step_scores = prev_step_scores.type(float_type).to(device)
self.last_state = None
def get_inputs(self) -> Gpt2BeamSearchInputs:
return Gpt2BeamSearchInputs(
self.input_ids,
self.position_ids,
self.attention_mask,
self.past,
self.beam_select_idx,
self.input_log_probs,
self.input_unfinished_sents,
self.prev_step_results,
self.prev_step_scores,
)
def update(self, output, step, device):
"""
Update the inputs for next inference.
"""
self.last_state = (
torch.from_numpy(output[0]).to(device)
if isinstance(output[0], numpy.ndarray)
else output[0].clone().detach().cpu()
)
self.input_ids = self.last_state.view(self.batch_size * self.beam_size, -1).to(device)
self.beam_select_idx = (
torch.from_numpy(output[-5]).to(device)
if isinstance(output[-5], numpy.ndarray)
else output[-5].clone().detach().to(device)
)
self.input_log_probs = (
torch.from_numpy(output[-4]).to(device)
if isinstance(output[-4], numpy.ndarray)
else output[-4].clone().detach().to(device)
)
self.input_unfinished_sents = (
torch.from_numpy(output[-3]).to(device)
if isinstance(output[-3], numpy.ndarray)
else output[-3].clone().detach().to(device)
)
self.prev_step_results = (
torch.from_numpy(output[-2]).to(device)
if isinstance(output[-2], numpy.ndarray)
else output[-2].clone().detach().to(device)
)
self.prev_step_scores = (
torch.from_numpy(output[-1]).to(device)
if isinstance(output[-1], numpy.ndarray)
else output[-1].clone().detach().to(device)
)
self.top_1_tokens = self.input_ids[0]
self.top_k_tokens = self.last_state
self.position_ids = (
torch.tensor([self.input_length + step - 1])
.unsqueeze(0)
.repeat(self.batch_size * self.beam_size, 1)
.to(device)
)
if self.attention_mask.size(0) != (self.batch_size * self.beam_size):
self.attention_mask = self.attention_mask.repeat(
self.batch_size * self.beam_size, 1
)
self.attention_mask = torch.cat(
[
self.attention_mask,
torch.ones([self.batch_size * self.beam_size, 1]).type_as(
self.attention_mask
),
],
1,
).to(device)
self.past = []
if isinstance(output[1], tuple): # past in torch output is tuple
self.past = list(output[1])
else:
for i in range(self.n_layer):
past_i = (
torch.from_numpy(output[i + 1])
if isinstance(output[i + 1], numpy.ndarray)
else output[i + 1].clone().detach()
)
self.past.append(past_i.to(device))
@staticmethod
def test_generation(session,
model,
device,
test_inputs,
precision=Precision.FLOAT32,
model_class="GPT2LMHeadModel_BeamSearchStep",
top_k=20,
top_k_no_order=True,
max_steps=24,
max_inputs=0,
verbose=False,
save_test_data=0,
save_test_data_dir="."):
"""
Test Generation using beam search to compare PyTorch and ONNX model.
It will print top 1 and top k errors on the given test inputs.
"""
print(
f"start test generation: (top_k={top_k} top_k_no_order={top_k_no_order} max_steps={max_steps} test_inputs={len(test_inputs)} max_inputs={max_inputs})"
)
n_layer = model.config.n_layer
n_head = model.config.n_head
n_embd = model.config.n_embd
beam_size = model.config.beam_size
eos_token_id = model.config.eos_token_id
test_data_saved = 0
is_float16 = precision == Precision.FLOAT16
# We will still use fp32 torch model as baseline when onnx model if fp16
model.eval().to(device)
# Allocate initial buffers for IO Binding of ONNX Runtimne. The buffer size will automatically increase later.
init_output_shapes = Gpt2BeamSearchHelper.get_output_shapes(
batch_size=4,
context_len=128,
past_sequence_length=128,
sequence_length=32,
beam_size=1,
step=0,
config=model.config,
model_class=model_class,
)
output_buffers = Gpt2BeamSearchHelper.get_output_buffers(
init_output_shapes,
device,
is_float16=is_float16,
)
baseline_name = "Torch"
treatment_name = "Quantized Onnx" if precision == Precision.INT8 else "Onnx"
torch_metric = Gpt2Metric(baseline_name, baseline_name, top_k)
onnx_metric = Gpt2Metric(treatment_name, baseline_name, top_k)
onnx_io_metric = Gpt2Metric(
treatment_name + " with IO Binding", baseline_name, top_k
)
for i, inputs in enumerate(test_inputs):
if max_inputs > 0 and i == max_inputs:
break
if i % 10 == 0:
print(f"{i}")
input_ids = inputs["input_ids"]
position_ids = inputs["position_ids"] if "position_ids" in inputs else None
attention_mask = (
inputs["attention_mask"] if "attention_mask" in inputs else None
)
beam_select_idx = (
inputs["beam_select_idx"] if "beam_select_idx" in inputs else None
)
input_log_probs = (
inputs["input_log_probs"] if "input_log_probs" in inputs else None
)
input_unfinished_sents = inputs["input_unfinished_sents"]
prev_step_results = inputs["input_ids"]
if "prev_step_scores" in inputs:
prev_step_scores = inputs["prev_step_scores"]
else:
prev_step_scores = torch.zeros([input_ids.shape[0], 1])
onnx_runner = Gpt2BeamSearchTester(
input_ids,
position_ids,
attention_mask,
beam_select_idx,
input_log_probs,
input_unfinished_sents,
prev_step_results,
prev_step_scores,
n_head,
n_embd,
n_layer,
beam_size,
device,
is_float16,
top_k,
not top_k_no_order,
)
onnx_io_runner = Gpt2BeamSearchTester(
input_ids,
position_ids,
attention_mask,
beam_select_idx,
input_log_probs,
input_unfinished_sents,
prev_step_results,
prev_step_scores,
n_head,
n_embd,
n_layer,
beam_size,
device,
is_float16,
top_k,
not top_k_no_order,
)
torch_runner = Gpt2BeamSearchTester(
input_ids,
position_ids,
attention_mask,
beam_select_idx,
input_log_probs,
input_unfinished_sents,
prev_step_results,
prev_step_scores,
n_head,
n_embd,
n_layer,
beam_size,
device,
False,
top_k,
not top_k_no_order,
) # Torch model baseline is fp32
batch_size = torch_runner.batch_size
onnx_metric.start_batch(batch_size)
onnx_io_metric.start_batch(batch_size)
context_len = list(onnx_runner.input_ids.size())[1]
with torch.no_grad():
done = torch.zeros(batch_size, dtype=torch.bool)
for step in range(max_steps):
print(f"Processing step: {step}")
seq_len = list(onnx_runner.input_ids.size())[1]
past_seq_len = list(onnx_runner.past[0].size())[3]
start_time = timeit.default_timer()
pytorch_output = Gpt2BeamSearchHelper.pytorch_inference(
model, torch_runner.get_inputs()
)
torch_metric.add_latency(
past_seq_len, timeit.default_timer() - start_time
)
torch_runner.update(pytorch_output, step, device)
(
onnx_output,
avg_latency_ms,
) = Gpt2BeamSearchHelper.onnxruntime_inference(
session, onnx_runner.get_inputs(), total_runs=1
)
onnx_metric.add_latency(past_seq_len, avg_latency_ms / 1000.0)
onnx_runner.update(onnx_output, step, device)
output_shapes = Gpt2BeamSearchHelper.get_output_shapes(
batch_size,
context_len,
past_seq_len,
seq_len,
beam_size,
step,
model.config,
model_class=model_class
)
Gpt2BeamSearchHelper.auto_increase_buffer_size(
output_buffers, output_shapes
)
(
onnx_io_output,
avg_latency_ms,
) = Gpt2BeamSearchHelper.onnxruntime_inference_with_binded_io(
session,
onnx_io_runner.get_inputs(),
output_buffers,
output_shapes,
total_runs=1,
return_numpy=False,
include_copy_output_latency=True,
)
onnx_io_metric.add_latency(past_seq_len, avg_latency_ms / 1000.0)
if test_data_saved < save_test_data:
onnx_io_runner.save_test_data(
session, onnx_io_output, save_test_data_dir, test_data_saved
)
test_data_saved += 1
onnx_io_runner.update(onnx_io_output, step, device)
done = done | (not onnx_runner.input_unfinished_sents.all())
if torch.all(done):
print("break at step: ", step)
break
print(f"Totally {step+1} steps run")
onnx_metric.end_batch()
onnx_io_metric.end_batch()
torch_metric.print()
onnx_metric.print()
onnx_io_metric.print()
print("\tONNX")
Gpt2BeamSearchTester.pprint_results(
onnx_runner.prev_step_results.view(batch_size * beam_size, -1),
onnx_runner.prev_step_scores.view(batch_size * beam_size, -1),
pad_token_id=eos_token_id,
eos_token_id=eos_token_id,
)
print("\tONNX with IO binding")
Gpt2BeamSearchTester.pprint_results(
onnx_io_runner.prev_step_results.view(batch_size * beam_size, -1),
onnx_io_runner.prev_step_scores.view(batch_size * beam_size, -1),
pad_token_id=eos_token_id,
eos_token_id=eos_token_id,
)
@staticmethod
def pprint_results(
output_ids,
output_scores,
pad_token_id=None,
eos_token_id=None,
):
"""
Print test generation results.
"""
if pad_token_id is None:
pad_token_id = 1
if eos_token_id is None:
eos_token_id = 1
if torch.is_tensor(output_ids):
output_ids = output_ids.cpu().numpy()
for i, sample in enumerate(output_ids):
for j, seq in enumerate(sample):
if isinstance(seq, numpy.ndarray) or isinstance(seq, list):
# remove left padding
for k, t in enumerate(seq):
if t != pad_token_id:
seq = seq[k:]
break
# remove EOS
for k, t in enumerate(seq):
if t == eos_token_id:
seq = seq[: k + 1]
break
print("-" * 40)
result = ",".join([str(token_id) for token_id in sample])
print(f">> Output {j + 1}: \t{[result]}")
else:
result = ",".join([str(token_id) for token_id in sample])
print(f">> Output {i}: \t{result}")
print(f">> Scores {i}: \t{output_scores[i]}")
break
print("=" * 80)