# Copyright (c) 2023 Graphcore Ltd. All rights reserved. import argparse import logging import os import queue import sys import tempfile import time from typing import Any import numpy as np import onnx import poprt from poprt import runtime from poprt.compiler import Compiler, CompilerOptions sys.path.append("..") import helper import onnxruntime np.random.seed(2023) def export_bert(dir): os.system( f"python -m transformers.onnx --model=bert-large-uncased {dir} --feature question-answering" ) os.system( f"python ../../examples/packed_bert_example/add_position_ids.py --input_model {dir}/model.onnx" ) class BertInputs(object): def __init__( self, input_ids, attention_mask, token_type_ids, position_ids, input_len ): self.input_ids = np.expand_dims(input_ids, axis=0) self.attention_mask = np.expand_dims(attention_mask, axis=0) self.token_type_ids = np.expand_dims(token_type_ids, axis=0) self.position_ids = np.expand_dims(position_ids, axis=0) self.input_len = [input_len] def gen_datasets(args, dtype: Any, random_data: bool = True): raw_inputs = [] if random_data: print( f"use random data,avg_seq_len: {args.avg_seq_len}, max_seq_len: {args.max_seq_len}" ) input_len = np.random.normal( args.avg_seq_len, args.avg_seq_len, size=args.dataset_size ).astype(np.int32) input_len = np.clip(input_len, 1, args.max_seq_len) for s_len in input_len: input_ids = np.random.randint(0, args.emb_size, (s_len)).astype(np.int32) attention_mask = np.ones(s_len).astype(np.int32) token_type_ids = np.random.randint(0, 1, (s_len)).astype(np.int32) position_ids = np.arange(s_len).astype(np.int32) feature = BertInputs( input_ids, attention_mask, token_type_ids, position_ids, s_len ) raw_inputs.append(feature) return raw_inputs else: from datasets import load_dataset from transformers import BertTokenizer tokenizer = BertTokenizer.from_pretrained('csarron/bert-base-uncased-squad-v1') raw_question = load_dataset("squad") text = [ "Where is london located", ] text = text * args.dataset_size for elem in text[: args.dataset_size]: feature = tokenizer( elem, return_tensors='np', max_length=args.max_seq_len, truncation=True, ) trans_feature = BertInputs( input_ids=np.array(feature['input_ids'][0], dtype), token_type_ids=np.array(feature['token_type_ids'][0], dtype), attention_mask=np.array(feature['attention_mask'][0], dtype), input_len=len(feature['input_ids'][0]), position_ids=np.arange(len(feature['input_ids'][0])).astype(dtype), ) raw_inputs.append(trans_feature) dataset_size = len(raw_question["validation"]) total_len = 0 for elem in raw_inputs: total_len += elem.input_len[0] avg_len = total_len / dataset_size print(f"validation dataset size: {dataset_size} avg_len: {avg_len}") return raw_inputs def run_bert_pack(args, popef, datasets, raw_model, out_names): config = runtime.PackRunnerConfig( timeout_microseconds=args.timeout_microseconds, max_valid_num=args.max_valid_num, dynamic_input_name=args.dynamic_input_name, ) # here we use first_fit algo config.algorithm = runtime.PackAlgorithm.first_fit config.enable_input_single_row_mode("attention_mask") pack_runner = runtime.Runner(popef, config) result_queue = queue.Queue() results = [] start_time = time.time() for i in range(args.dataset_size): feed_dicts = { "input_ids": datasets[i].input_ids, "attention_mask": datasets[i].attention_mask, "token_type_ids": datasets[i].token_type_ids, "position_ids": datasets[i].position_ids, } out_dict = { "start_logits": np.zeros([args.max_seq_len]).astype(np.float16), "end_logits": np.zeros([args.max_seq_len]).astype(np.float16), } future = pack_runner.execute_async(feed_dicts, out_dict) out_dict = {k: v[: datasets[i].input_len[0]] for k, v in out_dict.items()} result_queue.put((future, out_dict)) result_queue.put((None, None)) while True: future, out_dict = result_queue.get() if future == None: break future.wait() results.append(out_dict) end_time = time.time() tput = args.dataset_size / (end_time - start_time) latency_ms = (end_time - start_time) / args.dataset_size logging.info( f"Pack bert large Batch Size: {args.batch_size} Throughput: {tput} samples/s, Latency : {latency_ms * 1000} ms" ) if not args.random_data: ort_outputs = [] for raw_input in datasets[:30]: raw_input = raw_input.__dict__ raw_input.pop("input_len") raw_input.pop("position_ids") input = {} for k, v in raw_input.items(): input[k] = np.array(v).astype(np.int64) ort_session = onnxruntime.InferenceSession(raw_model) ort_output = ort_session.run(out_names, input) ort_outputs.append([np.array(elem).flatten() for elem in ort_output]) ipu_outputs = [list(e.values()) for e in results[:30]] ipu_outputs = np.concatenate( [np.concatenate((sl, el)) for sl, el in ipu_outputs] ) ort_outputs = np.concatenate( [np.concatenate((sl, el)) for sl, el in ort_outputs] ) helper.accuracy(ort_outputs, ipu_outputs) def test_func(args): batch_size = args.batch_size bps = args.bps sequence_len = args.max_seq_len input_shape = { 'input_ids': [batch_size, sequence_len], 'attention_mask': [batch_size, sequence_len], 'token_type_ids': [batch_size, sequence_len], 'position_ids': [batch_size, sequence_len], } with tempfile.TemporaryDirectory() as tmpdir: onnx_file = "model.onnx" bert_pos_path = "bert_base_squad_pos.onnx" export_bert(dir=tmpdir) raw_onnx_file = tmpdir + "/" + onnx_file logging.info( f"start benchmark for model: {raw_onnx_file},input_shape {input_shape}" ) bert_onnx = onnx.load(bert_pos_path) model = poprt.Converter( input_shape=input_shape, precision='fp16', enable_erf_gelu=True, used_passes=[ "insert_attention_mask", "pre_scale", "fused_attention", ], ).convert(bert_onnx) options = CompilerOptions() options.enable_prefetch_datastreams = True options.batches_per_step = bps options.partials_type = 'half' options.use_128bit_conv_unit_load = True options.enable_multi_stage_reduce = False options.enable_fast_reduce = True options.available_memory_proportion = 0.15 options.outline_threshold = 1e-7 out_names = [o.name for o in model.graph.output] popef = Compiler.compile(model, out_names, options) datasets = gen_datasets(args, np.int32, random_data=args.random_data) run_bert_pack(args, popef, datasets, raw_onnx_file, out_names) if __name__ == '__main__': logging.basicConfig(format='%(levelname)s:%(message)s', level=logging.INFO) parser = argparse.ArgumentParser(description='Inference bert with Poprt') parser.add_argument( '--batch_size', type=int, default=1, help='batch size of inference' ) parser.add_argument('--random_data', action='store_true', help='if use random data') parser.add_argument('--bps', type=int, default=100, help='batch per step') parser.add_argument('--dataset_size', type=int, default=1000, help='dataset size') parser.add_argument( '--dynamic_input_name', type=str, default="input_ids", help='input lenght that is dynamic', ) parser.add_argument( '--max_seq_len', type=int, default=384, help='max sequence length' ) parser.add_argument( '--avg_seq_len', type=int, default=192, help='average sequence length of input' ) parser.add_argument('--max_valid_num', type=int, default=64, help='max valid num') parser.add_argument( '--emb_size', type=int, default=30522, help='word embedding table size' ) parser.add_argument( '--timeout_microseconds', type=int, default=1500, help='longest time cpu pack wait', ) args = parser.parse_args() test_func(args)