Batchサイズの変更

複数の推論リクエストをグループ化してバッチ処理する事で、高いスループットを実現する事が可能です。グループ化された各推論リクエストに対して必要となるパラーメータは外部メモリへのアクセス無しで再利用されます。 Neuronコア Batchingによる最適化を行うためにはコンパイル時にバッチサイズを指定する必要があります。

Batchingにはレイテンシのコストを払うことになるため、最適なバッチサイズ（最適なレイテンシとスループットの組み合わせ）はユーザによるトライアルエラーにより決定されます。

Step 1. コンパイル用 Python スクリプトを作成

章5.1で作成したinfer_bert_cpu.pyの最後に以下の内容を追加し、 compile_bert_batch.py というファイル名のコンパイル用 Python スクリプトを作成します。 ここでは batch サイズ 6のモデルを用意します。

import torch.neuron

batch_size = 6

example_inputs_paraphrase = (
    torch.cat([paraphrase['input_ids']] * batch_size, 0),
    torch.cat([paraphrase['attention_mask']] * batch_size, 0),
    torch.cat([paraphrase['token_type_ids']] * batch_size, 0)
)

# Run torch.jit.trace to generate a TorchScript for running on CPU
model_cpu_batch = torch.jit.trace(model, example_inputs_paraphrase)

# Run torch.neuron.trace to generate a TorchScript that is optimized by AWS Neuron
model_neuron_batch = torch.neuron.trace(model, example_inputs_paraphrase)

## Save the batched model
model_cpu_batch.save('bert_cpu_b6.pt')
model_neuron_batch.save('bert_neuron_b6.pt')

Step 2. コンパイル用スクリプトを実行

コンパイルスクリプトを実行します。inf1.2xlarge では~4分程度かかります。

python compile_bert_batch.py

Step 3. 推論スクリプトの変更

以下の内容でinfer_bert_batch.py というファイル名の推論実行 Python スクリプトを作成します。 章5.4で作成したinfer_bert_perf2.pyとの違いは、batch_sizeの設定を1から6に変更した点、それに合わせて使用するモデルを変更した点の2点のみです。

import os
import time
import numpy as np
import torch
import torch_neuron
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from concurrent import futures

model_path = 'bert_neuron_b6.pt'

num_neuroncore = 4
num_thread = 2
batch_size = 6

throughput_time = 90
throughput_interval = 10
latency_window_size = 1000

# added for utilizing 4 neuron cores
os.environ['NEURON_RT_VISIBLE_CORES'] = '0-3'

# Build tokenizer
tokenizer = AutoTokenizer.from_pretrained("bert-base-cased-finetuned-mrpc")

# Load TorchScript back
model_list = [torch.jit.load(model_path) for _ in range(num_neuroncore)]
model_mp_list = []
for model in model_list:
    model_mp_list.extend(model for _ in range(num_thread))

# Setup some example inputs
sequence_0 = "The company HuggingFace is based in New York City"
sequence_1 = "Apples are especially bad for your health"
sequence_2 = "HuggingFace's headquarters are situated in Manhattan"

paraphrase = tokenizer.encode_plus(sequence_0, sequence_2, max_length=128, padding='max_length', truncation=True, return_tensors="pt")

# Convert example inputs to a format that is compatible with TorchScript tracing
example_inputs_paraphrase_tuple = (
    torch.cat([paraphrase['input_ids']] * batch_size, 0),
    torch.cat([paraphrase['attention_mask']] * batch_size, 0),
    torch.cat([paraphrase['token_type_ids']] * batch_size, 0)
)


live = True
num_infer = 0
latency_list = []

def one_thread(model, example_inputs_paraphrase_tuple):
    global latency_list
    global num_infer
    global live

    while True:
        start = time.time()
        paraphrase_classification_logits = model(*example_inputs_paraphrase_tuple)
        latency = time.time() - start
        latency_list.append(latency)
        num_infer += batch_size
        if not live:
            break

def current_performance():
    last_num_infer = num_infer
    for _ in range(throughput_time // throughput_interval):
        current_num_infer = num_infer
        throughput = (current_num_infer - last_num_infer) / throughput_interval
        p50 = 0.0
        p90 = 0.0
        p95 = 0.0
        if latency_list:
            p50 = np.percentile(latency_list[-latency_window_size:], 50)
            p90 = np.percentile(latency_list[-latency_window_size:], 90)
            p95 = np.percentile(latency_list[-latency_window_size:], 95)
        print('current throughput {}, latency p50={:.5f} p90={:.5f} p95={:.5f}'.format(throughput, p50, p90, p95))
        last_num_infer = current_num_infer
        time.sleep(throughput_interval)
    global live
    live = False

executor = futures.ThreadPoolExecutor(max_workers=len(model_mp_list)+1)
executor.submit(current_performance)

for model in model_mp_list:
   executor.submit(one_thread, model, example_inputs_paraphrase_tuple)

Step 4. 変更後の推論スクリプトを実行

推論実行スクリプトinfer_bert_batch.pyを実行します。

python infer_bert_batch.py

出力結果から、Batchサイズを大きくした場合、Neuronコア上での推論スループットは大きく向上している事が確認できます。

current throughput 0.0, latency p50=0.00000 p90=0.00000 p95=0.00000
current throughput 1090.8, latency p50=0.04418 p90=0.04425 p95=0.04427
current throughput 1084.8, latency p50=0.04418 p90=0.04425 p95=0.04427
current throughput 1087.2, latency p50=0.04418 p90=0.04425 p95=0.04426
current throughput 1087.2, latency p50=0.04418 p90=0.04425 p95=0.04427
current throughput 1087.2, latency p50=0.04418 p90=0.04425 p95=0.04427
current throughput 1086.0, latency p50=0.04418 p90=0.04424 p95=0.04426
current throughput 1086.0, latency p50=0.04418 p90=0.04425 p95=0.04426
current throughput 1092.0, latency p50=0.04385 p90=0.04396 p95=0.04401