mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-05-21 21:52:11 +00:00
86 lines
3.5 KiB
Python
86 lines
3.5 KiB
Python
#!/usr/bin/env python3
|
|
# Copyright (c) Microsoft Corporation. All rights reserved.
|
|
# Licensed under the MIT License.
|
|
|
|
import argparse
|
|
import subprocess
|
|
import sys
|
|
import os
|
|
import json
|
|
from collections import namedtuple
|
|
|
|
SCRIPT_DIR = os.path.realpath(os.path.dirname(__file__))
|
|
|
|
def parse_args():
|
|
parser = argparse.ArgumentParser(description="Runs BERT performance tests.")
|
|
parser.add_argument("--binary_dir", required=True,
|
|
help="Path to the ORT binary directory.")
|
|
parser.add_argument("--training_data_root", required=True,
|
|
help="Path to the training data root directory.")
|
|
parser.add_argument("--model_root", required=True,
|
|
help="Path to the model root directory.")
|
|
parser.add_argument("--gpu_sku", choices=['V100_16G', 'MI100_32G'], default='V100_16G', required=False,
|
|
help="GPU model (e.g. V100_16G, MI100_32G).")
|
|
return parser.parse_args()
|
|
|
|
# using the same params from "GitHub Master Merge Schedule" in OneNotes
|
|
def main():
|
|
args = parse_args()
|
|
|
|
Config = namedtuple('Config', ['use_mixed_precision', 'max_seq_length', 'batch_size', 'max_predictions_per_seq', 'expected_perf'])
|
|
configs = {}
|
|
configs['V100_16G'] = [
|
|
Config(True, 128, 76, 20, -1.0),
|
|
Config(True, 512, 11, 80, -1.0),
|
|
Config(False, 128, 39, 20, -1.0),
|
|
Config(False, 512, 6, 80, -1.0)
|
|
]
|
|
|
|
configs['MI100_32G'] = [
|
|
Config(True, 128, 128, 20, 240),
|
|
]
|
|
|
|
# run BERT training
|
|
for c in configs[args.gpu_sku]:
|
|
model = 'bert-large-uncased_L_24_H_1024_A_16_V_30528_S_512_Dp_0.1_optimized_layer_norm_opset12'
|
|
precision_prefix = ('fp16' if c.use_mixed_precision else 'fp32')
|
|
print("######## testing name - " + ('fp16-' if c.use_mixed_precision else 'fp32-') + str(c.max_seq_length) + " ##############")
|
|
cmds = [
|
|
os.path.join(args.binary_dir, "onnxruntime_training_bert"),
|
|
"--model_name", os.path.join(
|
|
args.model_root, "nv/bert-large/{}".format(model)),
|
|
"--train_data_dir", os.path.join(
|
|
args.training_data_root, str(c.max_seq_length), "books_wiki_en_corpus/train"),
|
|
"--test_data_dir", os.path.join(
|
|
args.training_data_root, str(c.max_seq_length), "books_wiki_en_corpus/test"),
|
|
"--train_batch_size", str(c.batch_size),
|
|
"--mode", "train",
|
|
"--max_seq_length", str(c.max_seq_length),
|
|
"--num_train_steps", "640",
|
|
"--display_loss_steps", "5",
|
|
"--optimizer", "Lamb",
|
|
"--learning_rate", "3e-3",
|
|
"--warmup_ratio", "0.2843",
|
|
"--warmup_mode", "Poly",
|
|
"--gradient_accumulation_steps", "1",
|
|
"--max_predictions_per_seq", str(c.max_predictions_per_seq),
|
|
"--lambda", "0",
|
|
"--use_nccl",
|
|
"--perf_output_dir", os.path.join(SCRIPT_DIR, "results"),
|
|
]
|
|
|
|
if c.use_mixed_precision:
|
|
cmds.append("--use_mixed_precision"),
|
|
cmds.append("--allreduce_in_fp16"),
|
|
|
|
subprocess.run(cmds).check_returncode()
|
|
if c.expected_perf > 0.0:
|
|
json_filename = 'onnxruntime_perf_metrics_{}.onnx_bert_{}_{}_Lamb.json'.format(model, precision_prefix, c.max_seq_length)
|
|
with open(os.path.join(SCRIPT_DIR, 'results', json_filename)) as json_file:
|
|
results = json.load(json_file)
|
|
assert(results['EndToEndThroughput'] > 0.98*c.expected_perf)
|
|
|
|
return 0
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|