はじめての自然言語処理

第25回 Transformer 系モデルの推論高速化の検証

オージス総研技術部データエンジニアリングセンター

鵜野和也

2023年2月21日

今回は Transformer 系のモデル、具体的には BERT, T5, GPT の推論を高速化してみます。高速化手法として FasterTransformer, Torch-TensorRT, AWS Neuron を用い、素の transfomers に比べ、どの程度速くなるか（ならないか）、利点・欠点を確認してみましょう。

1. はじめに

今回は Transformer 系のモデル、具体的には BERT, T5, GPT の推論を様々な技術を使って高速化してみます。高速化の元ネタは Hugging Face の transformers¹ 縛りとして、素の transformers で推論する場合に比べ、どの程度速くなるか（ならないか）見てみましょう。

推論を高速化する技術としては FasterTransfomer², Torch-TensorRT³, AWS Neuron(以後、単に Neuronと表記)⁴ を使います。それぞれ利点・欠点があるので、実際にモデルを変換、推論速度を計測して、そのあたりを確認して行きます。

さて、モデルの推論時は NVIDIA 製の GPU や AWS の Inferentia⁵ 等の特殊なデバイスを利用することが多いので、アプリケーションから利用する場合は Web API の形式で分離されていると色々と好都合ですよね。

ですので、今回は文字列を入力として受け取り推論結果を返す Web API を作って、そのスループットを評価することにしたいと思います。計測値は Tokenizer のエンコード/デコードやエンティティの切り出し処理等々を含んだ値になるので、高速化技術の純粋な比較ではないかもですが、出来るだけ条件は合わせていきたいと思います⁶。

Web API の構成としては、第14回で行ったように gunicorn⁷/Falcon⁸ を使いました。これをフロントエンドとして Tokenizer のエンコード/デコード等の前/後処理を行います。モデルの推論処理の配置は 2 パターンで以下のようなイメージになります。

素の transformers 及び Neuron の場合は推論もフロント側で処理⁹。
FasterTransformer, Torch-TensorRT はバックエンドに Triton Inference Server¹⁰ を立て、そちらで推論。

いろいろと登場する技術が多いので、一つ一つをあまり深くご説明はできませんが、良さげだと思ったものをより深く調べて頂くきっかけになればと思います。

では、ちょうど名前が出てきたので、まずは Triton Inference Server のご紹介です。

2. Triton Inference Server

Triton Inference Server (以下、Triton と表記）は NVIDIA の開発したオープンソースの推論サーバです。似たような立ち位置としては Google の Tensorflow Serving¹¹ がありますね。Triton の特徴を簡単にまとめると以下のようになります。

複数の深層学習フレームワーク（TensorRT, TensorFlow, PyTorch, ONNX, OpenVINO, Python, RAPIDS FIL, その他）に対応している。
多様なデバイス（NVIDIA GPU, x86 / ARM CPU, AWS Inferentia）に対応している。
多様な処理形態（リアルタイム、バッチ処理、音声／動画のストリーミング等）に対応している。
複数モデルや単一モデルの複数インスタンスの並行処理に対応している。
動的バッチングにより単発のリクエストを内部的にバッチにまとめて処理できる。
ステートフルなモデルに対応できる（同じシーケンスに属するリクエストを同一のモデルインスタンスにディスパッチする）。
HTTP/GRPC に対応。

他にもいろいろとありますが書いているとキリがないので、 Triton にモデルをデプロイする際にキモになる Model Repository について説明しておきます。

Model Repository

Model Repository¹² は名前そのままで Triton にデプロイするモデルを格納するディレクトリ（及びその構造）のことで以下のような構成になります。

model-repository-path :
Model Repository のパス。ローカルストレージの他に GCS や S3 を使用することもできます。
model-name :
デプロイするモデルの名前です。この名前が推論時にモデルを指定する際に使用されます。
config.pbtx :
モデルの設定ファイルです。このファイルに使用するバックエンド(TensorRT, Tensorflow, ONIX 等）、入出力のシェイプ/データ型、動的バッチングやマルチインスタンスの構成等を記述します。使用するバックエンドによっては必須項目を自動生成することもできるようです。
output-labels-file :
執筆時点で初めて見ました。私は使ったことありません。執筆時点のドキュメント¹²にも特に記載ありませんね。。。
version :
デプロイするモデルのバージョンを示す数値を名前とするディレクトリです。Triton はモデルに対し複数バージョンを保持することができ、同時あるいは選択的にデプロイ可能です( config.pbtxt で設定します)。
model-definition-file :
デプロイするモデルの実体（Tensorflow なら SavedModel に相当）です。実際のファイル名は使用するバックエンドによって異なるので、詳しくはこちら¹³を参照してください。

Triton の説明はこれぐらいにして、次は今回使った高速化手法について説明していきます。まずは FasterTransformer から見ていきましょう。

3. FasterTransformer

FasterTransformer² はNVIDIA が開発した高度に最適化された Transformer 系モデルの実装です。

後述する TensorRT や Neuron とは異なり、FasterTransformer はあくまで最適化済みの実装なので、Transformer なら何でも高速化できるわけではなく、BERT, T5, GPT 等のサポート済みのモデル¹⁴のみが対象です。こう書くと TensorRT の方が自由度があって良さげなのですが、大規模なモデルや T5, GPT のような自己回帰によるテキスト生成モデルでは、実際に TensorRT で最適化して推論しようとすると、いろいろ複雑で面倒です¹⁵。

その点、 FasterTransformer の場合は、モデルの学習済みパラメータを FasterTransformer の形式に変換する必要こそありますが、モデルの最適化というステップが存在しないので（というか既に最適化済みみたいなものなので）、学習済みパラメータをロードすればきっちり動きます。

FasterTransformer については、本記事では「速くなればOK」のスタンスなので「高速化されたモデルの実装です」以上にあまり説明することもありません。興味のある人は github リポジトリ²の “docs” あたりを見てみて下さい。ただ BERT に関しては少し補足しておきます。

FasterTransformer の BERT

T5 や GPT 等のテキスト生成モデルでは入出力ともにトークン ID の系列になります。ですが BERT ではタスクによって出力の形式は様々です。 FasterTransformer はオプティマイザではなくモデルの実装なので、この点についてどうするのか？と思ったのですが、かなり割り切った対応でした。以下は transformers の BertForTokenClassification(トークン単位の分類を行う際に使用するクラス) の構造ですが、FasterTransformer がカバーしてくれるのは橙の BERT 本体部分のみです。

この部分だけ着目すれば、どのような使い方をするにせよ入出力の形式は [batch_size, sequence_length, hidden_size] になりますが、後は自前で対応する必要があるので、かなり扱いづらいと感じました。「せめて BertModel 相当の範囲は対応しといてよ。」と思いながらも試すだけ試したので結果は後述しますね。。。

また、FasterTransformer を Triton にデプロイしようと思うと FasterTransformer Backend のインストールが必要です。

FasterTransformer Backend

FasterTransformer は FasterTransformer Backend¹⁶ を使用することで Triton にデプロイすることが可能です。

FasterTransformer Backend のドキュメント¹⁶には “Note that this is a research and prototyping tool, not a formal product or maintained framework. ” と記されているのですが、巨大な GPT-3 モデルを複数台の GPU サーバ / Triton に分散配置して連動して動かすような派手なことをしなければ、大丈夫な範囲では大丈夫そうです（大丈夫じゃなかった点は分かった範囲で後述します）。

個人的に FasterTransformer Backend で嬉しいところは T5, GPT のような自己回帰によるテキスト生成モデルを Triton にデプロイできるところです。今までは、第14回のような構成にするか、Google の T5 実装から生成した SavedModel の形式¹⁷で Triton にデプロイするかしかなかったのですが、ようやく transformers で学習した T5, GPT を Triton にデプロイすることができました！（というか元々は「transformers で学習した T5, GPT を Triton にデプロイ出来ないのか？」で調べていて、FasterTransformer に行き着いたので高速化がオマケ要素だったりします。）

それでは、次に Torch-TensorRT の話をしましょう。

4. Torch-TensorRT

まず TensorRT¹⁸ は NVIDIA の開発した深層学習モデルの推論に特化した高度なオプティマイザとランタイムです。そして Torch-TensorRT は TensorRT を用いた PyTorch/TorchScript/FX のコンパイラという位置づけになります。

また、Torch-TensorRT は実行時に変換を行う JIT コンパイラではなく、予めコンパイルした結果を保存しておくことが出来ます。

ちなみに Tensorflow の場合、 Triton は “TensorFlow with TensorRT Optimization(TF-TRT)” に対応しており、Tensorflow の SavedModel を TensorRT で最適化できます。ですが、こちらは最初の推論リクエスト時に長時間のコンパイルが必要になるので、それに比べると Torch-TensorRT の方が使い勝手が良さそうです。

ただし、Torch-TensorRT を使うときは以下の点を抑えておいて下さい。

コンパイルする環境の GPU は推論時の GPU に合わせておく必要がある。
コンパイル時と推論時の環境で TensorRT のバージョンを合わせておく必要がある。

加えて、今回はコンパイル対象とするモデルの入力シェイプをバッチ軸含めて固定としました。 Torch-TensorRT は 1.3.0 で動的バッチに対応¹⁹したようなのですが、Triton にバンドルされている Torch-TensorRT は、最新のコンテナイメージである 22.12 でも 1.1.0a0 でした²⁰。Torch-TensorRT を入れ替えて Triton をリビルドするのは大変そうだったので今回は諦めました。。。

また、Torch-TensorRT でコンパイルしたモデルは TorchScript と同様の方法で Triton にデプロイ可能です。

続いて AWS の Neuron です。

5. Neuron

Neuron は AWS Inferentia²¹ や Trainium²² デバイスを用いて深層学習を高速化する SDK になります。コンパイラやドライバが同梱されており、Tensorflow や PyTorch 等の一般的なフレームワークから利用できます。

役割分担としては Inferentia が推論用、Trainium が学習用のデバイスということのようです。今回は「transformers で学習した PyTorch モデルの推論を高速化したい」というシナリオなので Inferentia を搭載した Inf1 インスタンスを用います。

最初は試すものが増えすぎると大変なのでスルーしようかと思ったのですが、AWS に同等の GPU ベースの Amazon EC2 インスタンスに比べて、スループットが最大 2.3 倍、推論あたりのコストが最大 70% 削減されます。²¹と言われると試したくなってしまいます。

ただ、ドキュメント²³を見るとなんでもかんでも高速化できるわけではなく、BERT のような Encoder-Only には対応していて、GPT や T5 のような Decoder-Only, Encoder-Decoder には対応できていないようです。ですが、さらに調べていると「T5 を Neuron で動かした」という人がいた²⁴のでダメ元で試してみることにしました。

あと、この issue²⁵ を見ると入力シェイプが固定でないとダメみたいですね。バッチ軸は可変にできなくもなさそうなのですが、今回は Triton へのデプロイを諦めたので（Web API で 1 件づつ飛んでくるリクエストをバッチにまとめるのは Triton の動的バッチングに頼っているので）、batch_size = 1 固定で実験しました。

さて、ここからは実際に実験した環境とその結果をご紹介していきます。今回は Colab で動かすという訳にもいかず、コードも結構膨らんでしまったので環境のセットアップとモデル変換のところを中心にかいつまんでご紹介していきます。

6. 実験環境の物理構成

今回は AWS の ap-northeast-1 リージョンで以下のように環境を構築しました。

上段:
FasterTransformer, Torch-TensorRT, TorchScript のモデルは g4dn.xlarge 上の Triton にデプロイして、gunicorn/Falcon を前段に配置しました。。
中段/下段:
transformers / Aws Neuron のモデルは gunicorn/Falcon の worker がモデルを保持する形です。ただし transformers は g4dn.xlarge , Neuron は inf1.xlarge のインスタンスです。
左側:
Web API に対する負荷の生成は t3.medium の locust で行いました。

上図の白枠がコンテナに相当します。A1, A2, B, C は便宜的に付けたコンテナの識別子です。

7. 実験環境のソフトウェア構成

基本的に Docker を使って構成しました。各コンテナで使用したソフトウェアのバージョンは Dockerfile の形で後述しますが、ホストVM には以下の AMI とドライバをインストールしています。

AMI : amzn2-ami-kernel-5.10-hvm-2.0.20221103.3-x86_64-gp2 (上図の全VMで共通)
- 執筆時点で既にこの AMI を選択出来なくなっていますが、同系統の Amazon Linux 2 AMI (HVM) Kernel 5.10 なら大丈夫かと思います。
NVIDIA Driver : 515.65.01 (上図の上段、中段のみ）
Neuron Driver : 2.3.26 (上図の下段のみ）

上段、中段の NVIDIA Driver のインストール、下段の VM に対する Neuron Driver のインストール手順を以下に示しておきます。

7.1 NVIDIA Driver のインストール

VM の起動

前述の AMI を選択し、マシンタイプ g4dn.xlarge で起動します。

NVIDIA Driver のインストール

以下のようにして NVIDIA ドライバをインストールしました。

筆者が作業したときにはこの問題²⁶にハマったので gcc のバージョンに 10.x を指定しています。

$ BASE_URL=https://us.download.nvidia.com/tesla
$ DRIVER_VERSION=515.65.01
$ curl -fSsl -O $BASE_URL/$DRIVER_VERSION/NVIDIA-Linux-x86_64-$DRIVER_VERSION.run
$ sudo yum install -y gcc kernel-devel-$(uname -r) kernel-headers-$(uname -r)
$ sudo CC=/usr/bin/gcc10-gcc sh NVIDIA-Linux-x86_64-$DRIVER_VERSION.run

ドライバをインストール中に X library path 及び X module path に関する警告と 32bit バージョンのインストールの確認がありますが前者は OK、後者は NO で大丈夫です。

NVIDIA Docker のインストール

後は NVIDIA Docker もインストールしておきましょう。

$ sudo yum install -y docker
$ sudo systemctl start docker
$ sudo usermod -a -G docker ssm-user
$ sudo systemctl enable docker
$ distribution=$(. /etc/os-release;echo $ID$VERSION_ID) && curl -s -L https://nvidia.github.io/nvidia-docker/$distribution/nvidia-docker.repo | sudo tee /etc/yum.repos.d/nvidia-docker.repo
$ sudo yum install nvidia-docker2
$ sudo yum clean expire-cache
$ sudo systemctl restart docker
$ sudo docker run --rm --gpus all nvidia/cuda:11.6.2-base-ubuntu20.04 nvidia-smi

7.2 Neuron Driver のインストール

VM の起動

前述の AMI を選択し、マシンタイプ inf1.xlarge (※モデルのコンパイル処理のみ inf1.2xlarge で作業しました²⁷）で起動します。この際、VM の IAM Role に AmazonEC2ContainerRegistryPowerUser を追加しました²⁸。

Neuron Driver のインストール

以下のようにして Neuron ドライバをインストールしました。

$ sudo tee /etc/yum.repos.d/neuron.repo > /dev/null <<EOF
[neuron]
name=Neuron YUM Repository
baseurl=https://yum.repos.neuron.amazonaws.com
enabled=1
metadata_expire=0
EOF

$ sudo rpm --import https://yum.repos.neuron.amazonaws.com/GPG-PUB-KEY-AMAZON-AWS-NEURON.PUB

$ sudo yum update -y
$ sudo yum install kernel-devel-$(uname -r) kernel-headers-$(uname -r) -y
$ sudo yum install aws-neuron-dkms -y

$ ls /dev/neuron*
/dev/neuron0

Docker のインストール

以下のようにして Docker をインストールします。コンテナ内から Neuron を使う場合でも普通にインストールすれば大丈夫でした。

$ sudo yum install -y docker
$ sudo systemctl enable docker
$ sudo systemctl start docker
$ sudo usermod -a -G docker ssm-user # Please logout and login now.

続いて各コンテナの構成です。

以降で Dockerfile を複数作成しますが、全て /home/ssm-user/build に保存する想定として下さい。

7.3 コンテナ A1 の構成

このコンテナは Torch-TensorRT, TorchScript, FasterTransformer に変換したモデルを推論する場合のフロントエンド側になります。

このコンテナには Triton のクライアントが必要なので nvcr.io/nvidia/tritonserver:22.07-py3-sdk を親イメージとしています。また今回、Triton はコンテナイメージのバージョン 22.07 を用いました²⁹。

DockerFile.a1 は以下のとおりです。

FROM nvcr.io/nvidia/tritonserver:22.07-py3-sdk

RUN apt-get clean && apt-get update && \
    apt-get install -y mecab mecab-ipadic-utf8 && \
    curl https://bootstrap.pypa.io/pip/3.6/get-pip.py -o /tmp/get-pip.py && \
    python /tmp/get-pip.py && \
    pip install torch==1.12.1 transformers==4.23.1 \
                mecab-python3==1.0.5 fugashi==1.1.0 ipadic==1.0.0 \
                sentencepiece==0.1.97 \
                scikit-learn==1.1.3 \
                falcon==3.1.0 gunicorn==20.1.0

コンテナイメージを container-a1 としてビルドしておきます。

$ cd /home/ssm-user/build
$ docker build -t container-a1 -f DockerFile.a1 .

7.4 コンテナ A2 の構成

このコンテナは Torch-TensorRT, TorchScript, FasterTransformer に変換したモデルを推論する場合のバックエンド側の Triton になります。

NVIDIA から配布されている Triton の Docker イメージには FasterTransformer が組み込まれていないので追加してビルドし直します。ベースとする Triton のコンテナイメージは 22.07 を用いました。

$ pwd 
/home/ssm-user
$ git clone https://github.com/triton-inference-server/fastertransformer_backend
$ cd fastertransformer_backend
$ git checkout 225b57898b830a
$ docker build --build-arg TRITON_VERSION=22.07 -t container-a2 -f docker/Dockerfile .

この手順ですと FasterTransformer はビルド時点の main ブランチの最新になりますが、たしかこれ³⁰だったと思います。ビルド対象の FasterTransformer を固定したいときは

$ sed -i 's#main#release/v5.2_tag#' CMakeLists.txt

のようにして CMakeLists.txt の main を参照している部分をタグに置き換えれば大丈夫そうです。

7.5 コンテナ B の構成

このコンテナは transformers モデルをそのまま使って推論する場合の構成になります。

Dockerfile は。。。確認したらコンテナ A1 の物と同じでした。この構成では Triton のクライアントは不要なのですが、面倒なので同じ Dockerfile で作業してしまったみたいです。

後で混乱しそうなので先程の container-a1 のイメージに対して container-b として別名を付けておきました。

$ docker tag container-a1 container-b

7.6 コンテナ C の構成

このコンテナはコンテナ B とほぼ同じですが、Inf1 インスタンス上で稼働させることを前提とし、transformers のモデルを Neuron でコンパイルしたモデルで推論する場合の構成になります。

今回は Neuron SDK がインストール済みのイメージ(1.10.2-neuron-py37-sdk1.19.0-ubuntu18.04 )を使いました。まずはイメージを pull 出来るように ECR の認証を通します。

$ aws ecr get-login-password --region  ap-northeast-1 | docker login --username AWS --password-stdin 763104351884.dkr.ecr.ap-northeast-1.amazonaws.com

Dockerfile.c は以下のようにしました。

FROM 763104351884.dkr.ecr.ap-northeast-1.amazonaws.com/pytorch-inference-neuron:1.10.2-neuron-py37-sdk1.19.0-ubuntu18.04
RUN apt-get clean && apt-get update || echo "WARN : Ignoring Certificate verification failed"
WORKDIR /opt
RUN apt-get install -y mecab mecab-ipadic-utf8 python-mecab libmecab-dev pciutils
RUN pip install transformers==4.20.1 sentencepiece==0.1.97 mecab-python3==1.0.5 fugashi==1.1.0 ipadic==1.0.0
RUN pip install falcon==3.1.0 gunicorn==20.1.0

後はビルドすればOKです。

$ cd /home/ssm-user/build
$ docker build -t container-c -f Dockerfile.c .

ここからは transformers のモデルを変換していきましょう。今回対象とするのは以下の 3 つのモデルです。

BERT
- 第12回の UD_Japanese-GSD v2.6-NE データセットで学習した固有表現抽出モデル。
- 事前学習モデル: cl-tohoku/bert-base-japanese-whole-word-masking
- 以後、bert-ner と表記
GPT
- 第7回のやさしい日本語コーパスを加工したデータセットで学習したテキスト変換モデル。
- 事前学習モデル: rinna/japanese-gpt2-medium
- 以後、gpt2-medium と表記
T5
- 第7回のやさしい日本語コーパスを加工したデータセットで学習したテキスト変換モデル。
- 事前学習モデル: megagonlabs/t5-base-japanese-web
- ただし Neuron 向けのみ T5 1.0 ベースの sonoisa/t5-base-japanese を使用
- 以後、t5-base, t5-1.0-base と表記

まずは FasterTransformer からです。

8. FasterTransfomer への変換

ここからは g4dn.xlarge インスタンスの /home/ssm-user/hf-models に transformers の学習済みモデルが配置されている前提で記述します。

$ pwd 
/home/ssm-user
$ ls
fastertransformer_backend  hf-models
$ ls hf-models
bert-ner  gpt2-medium  t5-base

各モデルの内容物はこんな感じです。

hf-models
  + bert-ner
  |   +  all_results.json
  |   +  config.json
  |   +  eval_results.json
  |   +  predictions.txt
  |   +  predict_results.json
  |   +  pytorch_model.bin
  |   +  special_tokens_map.json
  |   +  tokenizer_config.json
  |   +  trainer_state.json
  |   +  training_args.bin
  |   +  train_results.json
  |   +  vocab.txt
  + gpt2-medium  
  |   +  config.json
  |   +  pytorch_model.bin
  |   +  special_tokens_map.json
  |   +  spiece.model
  |   +  tokenizer_config.json
  |   +  tokenizer.json
  |   +  trainer_state.json
  |   +  training_args.bin
  |   +  train_results.txt
  + t5-base
  |   +  config.json
  |   +  pytorch_model.bin
  |   +  rng_state.pth
  |   +  scheduler.pt
  |   +  special_tokens_map.json
  |   +  spiece.model
  |   +  tokenizer_config.json
  |   +  tokenizer.json
  |   +  trainer_state_all.json
  |   +  trainer_state.json
  |   +  training_args.bin

BERT の変換

まずは BERT を変換します。transformers からの変換スクリプトは FasterTransformer に同梱されているので、それを利用します。

今見ると FasterTransformer に f27f8bb0786b5a8a69 を使ってますね。。。これは特に意味がなくて試行錯誤の名残です。いろいろありすぎて、きちんとバージョンを合わせて作業できなかっただけで。今回の実験結果はこの commit で変換処理を行っ（てしまっ）たので、そのままを記載してます。気になる人は release/v5.2_tag とかで実行してもらえれば良いでしょう。

$ cd /home/ssm-user
$ docker run --gpus all --rm -v `pwd`:/work container-a2 \
bash -c \
"git clone https://github.com/NVIDIA/FasterTransformer.git &&"\
"cd FasterTransformer && git checkout f27f8bb0786b5a8a69 &&"\
"cd /work &&"\
"python3 /workspace/FasterTransformer/examples/pytorch/bert/utils/huggingface_bert_convert.py "\
" -in_file ./hf-models/bert-ner "\
" -saved_dir ./ft_bert "\
" -infer_tensor_para_size 1 "\
" -weight_data_type fp16"
...
=============== Argument ===============
saved_dir: ./ft_bert
in_file: ./hf-models/bert-ner
training_tensor_para_size: 1
infer_tensor_para_size: 1
processes: 4
weight_data_type: fp16
========================================
[WARNING] cannot convert key 'embeddings.LayerNorm.weight'
[WARNING] cannot convert key 'embeddings.word_embeddings.weight'
[WARNING] cannot convert key 'embeddings.LayerNorm.bias'
[WARNING] cannot convert key 'embeddings.position_embeddings.weight'
[WARNING] cannot convert key 'pooler.dense.weight'
[WARNING] cannot convert key 'embeddings.token_type_embeddings.weight'
[WARNING] cannot convert key 'pooler.dense.bias'
$ ls ft_bert/
1-gpu

なにやら警告がでていますが、とりあえすモデルの変換ができました。中身を確認して見ましょう。

$ ls ft_bert/1-gpu
config.ini                                                    model.encoder.layer.4.attention.output.dense.bias.bin
model.encoder.layer.0.attention.output.dense.bias.bin         model.encoder.layer.4.attention.output.dense.weight.0.bin
model.encoder.layer.0.attention.output.dense.weight.0.bin     model.encoder.layer.4.attention.output.LayerNorm.bias.bin
...

config.ini と各レイヤのパラメータに対応するバイナリが沢山あります。デプロイするときはこの 1-gpu ディレクトリを Triton のモデルリポジトリに以下のようにはめ込みます。

+ models/
    + ft_bert/
        + config.pbtxt
        + 1/
          + 1-gpu/

config.pbtxt は以下のとおりです。

backend: "fastertransformer"
max_batch_size: 16
input [
  {
    name: "input_hidden_state"
    data_type: TYPE_FP16
    dims: [ -1, -1 ]
  },
  {
    name: "sequence_lengths"
    data_type: TYPE_INT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
  }
]
output [
  {
    name: "output_hidden_state"
    data_type: TYPE_FP16
    dims: [ -1, -1 ]
  }
]
instance_group [
  {
    count: 1
    kind : KIND_CPU
  }
]
parameters {
  key: "tensor_para_size"
  value: {
    string_value: "1"
  }
}
parameters {
  key: "pipeline_para_size"
  value: {
    string_value: "1"
  }
}
parameters {
  key: "data_type"
  value: {
    string_value: "fp16"
  }
}
parameters {
  key: "enable_custom_all_reduce"
  value: {
    string_value: "0"
  }
}
parameters {
  key: "model_type"
  value: {
    string_value: "bert"
  }
}
parameters {
  key: "int8_mode"
  value: {
    string_value: "0"
  }
}
parameters {
  key: "is_sparse"
  value: {
    string_value: "0"
  }
}
parameters {
  key: "is_remove_padding"
  value: {
    string_value: "1"
  }
}
version_policy: { specific { versions : [1] }}
dynamic_batching {}

GPT の変換

GPT の変換も同様です。

$ cd /home/ssm-user
$ docker run --gpus all --rm -v `pwd`:/work container-a2 \
bash -c \
"git clone https://github.com/NVIDIA/FasterTransformer.git &&"\
"cd FasterTransformer && git checkout f27f8bb0786b5a8a69 &&"\
"cd /work &&"\
"python3 /workspace/FasterTransformer/examples/pytorch/gpt/utils/huggingface_gpt_convert.py "\
" -in_file ./hf-models/gpt2-medium "\
" -saved_dir ./ft_gpt2 "\
" -infer_gpu_num 1 "\
" -weight_data_type fp16"
...
=============== Argument ===============
saved_dir: ./ft_gpt2
in_file: ./hf-models/gpt2-medium
trained_gpu_num: 1
infer_gpu_num: 1
processes: 4
weight_data_type: fp16
========================================
$ ls ft_gpt2/
1-gpu

モデルリポジトリへの配置は BERT と同様です。

config.pbtxt は以下のようになります。

backend: "fastertransformer"
max_batch_size: 16

model_transaction_policy {
  decoupled: False
}

input [
  {
    name: "input_ids"
    data_type: TYPE_UINT32
    dims: [ -1 ]
  },
  {
    name: "input_lengths"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
  },
  {
    name: "request_output_len"
    data_type: TYPE_UINT32
    dims: [ -1 ]
  },
  {
    name: "runtime_top_k"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "runtime_top_p"
    data_type: TYPE_FP32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "beam_search_diversity_rate"
    data_type: TYPE_FP32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "temperature"
    data_type: TYPE_FP32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "len_penalty"
    data_type: TYPE_FP32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "repetition_penalty"
    data_type: TYPE_FP32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "random_seed"
    data_type: TYPE_UINT64
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "is_return_log_probs"
    data_type: TYPE_BOOL
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "beam_width"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "start_id"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "end_id"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "stop_words_list"
    data_type: TYPE_INT32
    dims: [ 2, -1 ]
    optional: true
  },
  {
    name: "bad_words_list"
    data_type: TYPE_INT32
    dims: [ 2, -1 ]
    optional: true
  },
  {
    name: "prompt_learning_task_name_ids"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "request_prompt_embedding"
    data_type: TYPE_FP16
    dims: [ -1, -1 ]
    optional: true
  },
  {
    name: "request_prompt_lengths"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "request_prompt_type"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  }
]
output [
  {
    name: "output_ids"
    data_type: TYPE_UINT32
    dims: [ -1, -1 ]
  },
  {
    name: "sequence_length"
    data_type: TYPE_UINT32
    dims: [ -1 ]
  },
  {
    name: "cum_log_probs"
    data_type: TYPE_FP32
    dims: [ -1 ]
  },
  {
    name: "output_log_probs"
    data_type: TYPE_FP32
    dims: [ -1, -1 ]
  }
]
instance_group [
  {
    count: 1
    kind : KIND_CPU
  }
]
parameters {
  key: "tensor_para_size"
  value: {
    string_value: "1"
  }
}
parameters {
  key: "pipeline_para_size"
  value: {
    string_value: "1"
  }
}
parameters {
  key: "data_type"
  value: {
    string_value: "fp16"
  }
}
parameters {
  key: "model_type"
  value: {
    string_value: "GPT"
  }
}
parameters {
  key: "int8_mode"
  value: {
    string_value: "0"
  }
}
parameters {
  key: "enable_custom_all_reduce"
  value: {
    string_value: "0"
  }
}
version_policy: { specific { versions : [1] }}
dynamic_batching {}

T5 の変換

T5 のモデルも変換しておきましょう。

$ cd /home/ssm-user
$ docker run --gpus all --rm -v `pwd`:/work container-a2 \
bash -c \
"git clone https://github.com/NVIDIA/FasterTransformer.git &&"\
"cd FasterTransformer && git checkout f27f8bb0786b5a8a69 &&"\
"cd /work &&"\
"python3 /workspace/FasterTransformer/examples/pytorch/t5/utils/huggingface_t5_ckpt_convert.py "\
" -in_file ./hf-models/t5-base "\
" -saved_dir ./ft_t5 "\
" -inference_tensor_para_size 1 "\
" -weight_data_type fp16"
...
2023-01-19 02:02:24,703 __main__ [INFO]
=============== Argument ===============
2023-01-19 02:02:24,703 __main__ [INFO] saved_dir: ./ft_t5
2023-01-19 02:02:24,703 __main__ [INFO] in_file: ./hf-models/t5-base
2023-01-19 02:02:24,703 __main__ [INFO] inference_tensor_para_size: 1
2023-01-19 02:02:24,703 __main__ [INFO] weight_data_type: fp16
2023-01-19 02:02:24,703 __main__ [INFO] verbose: False
2023-01-19 02:02:24,703 __main__ [INFO] ========================================
2023-01-19 02:02:34,043 __main__ [WARNING] Not save encoder.embed_tokens.weight, using shared.weight directly.
2023-01-19 02:02:36,837 __main__ [WARNING] Not save decoder.embed_tokens.weight, using shared.weight directly.
2023-01-19 02:02:40,121 __main__ [INFO] Spend 0:00:15.418305 (h:m:s) to convert the model
$ ls ft_t5/
1-gpu

config.pbtxt は以下のようになります。

backend: "fastertransformer"
max_batch_size: 16
input [
  {
    name: "input_ids"
    data_type: TYPE_UINT32
    dims: [ -1 ]
  },
  {
    name: "sequence_length"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
  },
  {
    name: "runtime_top_k"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "runtime_top_p"
    data_type: TYPE_FP32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "beam_search_diversity_rate"
    data_type: TYPE_FP32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "temperature"
    data_type: TYPE_FP32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "len_penalty"
    data_type: TYPE_FP32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "repetition_penalty"
    data_type: TYPE_FP32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "random_seed"
    data_type: TYPE_UINT64
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "is_return_log_probs"
    data_type: TYPE_BOOL
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "max_output_len"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
  },
  {
    name: "beam_width"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "start_id"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "end_id"
    data_type: TYPE_UINT32
    dims: [ 1 ]
    reshape: { shape: [ ] }
    optional: true
  },
  {
    name: "bad_words_list"
    data_type: TYPE_INT32
    dims: [ 2, -1 ]
    optional: true
  },
  {
    name: "stop_words_list"
    data_type: TYPE_INT32
    dims: [ 2, -1 ]
    optional: true
  }
]
output [
  {
    name: "output_ids"
    data_type: TYPE_UINT32
    dims: [ -1, -1 ]
  },
  {
    name: "sequence_length"
    data_type: TYPE_UINT32
    dims: [ -1 ]
  },
  {
    name: "cum_log_probs"
    data_type: TYPE_FP32
    dims: [ -1 ]
  },
  {
    name: "output_log_probs"
    data_type: TYPE_FP32
    dims: [ -1, -1 ]
  }
]
instance_group [
  {
    count: 1
    kind : KIND_CPU
  }
]
parameters {
  key: "tensor_para_size"
  value: {
    string_value: "1"
  }
}
parameters {
  key: "pipeline_para_size"
  value: {
    string_value: "1"
  }
}
parameters {
  key: "data_type"
  value: {
    string_value: "fp16"
  }
}
parameters {
  key: "enable_custom_all_reduce"
  value: {
    string_value: "0"
  }
}
parameters {
  key: "model_type"
  value: {
    string_value: "T5"
  }
}

version_policy: { specific { versions : [1] }}
dynamic_batching {}

では続いて Torch-TensorRT です。

9. Torch-TensorRT への変換

作業環境は FasterTransformer と同様で g4dn.xlarge インスタンスで transformers のモデルが /home/ssm-user/hf-models にある前提です。

変換処理で新たな環境が必要になったので Dockerfile.d として記述しました。

FROM pytorch/pytorch:1.12.1-cuda11.3-cudnn8-runtime
RUN apt-get clean && apt-get update && \
    apt-get install -y mecab mecab-ipadic-utf8 python-mecab libmecab-dev
RUN pip install transformers==4.23.1 sentencepiece==0.1.97 \
                mecab-python3==1.0.5 fugashi==1.1.0 ipadic==1.0.0
RUN pip install nvidia-pyindex==1.0.9 && \
    pip install nvidia-tensorrt==8.4.1.5 && \
    pip install torch-tensorrt==1.2.0 --find-links https://github.com/pytorch/TensorRT/releases/expanded_assets/v1.2.0

最後の RUN コマンドで Torch-TensoRT 関係をインストールしています。注意点としては TensorRT は変換時と推論時でバージョンを揃える必要があるので、 Triton の 22.07 に合わせて TensorRT の 8.4.1.x 系とそれに対応した Torch-TensorRT 1.2.0 を入れました。

コンテナイメージを container-d としてビルドしておきます。

$ cd /home/ssm-user/build
$ docker build -t container-d -f DockerFile.d .

変換するモデルですが今回は BERT のみ対象としました。GPT と T5 に関しては変換、Triton にデプロイできたとしても、自己回帰でテキスト生成するにはトークンの生成毎に Triton のエンドポイントを叩くしか方法がなさそうだったので諦めました。

BERT の変換用スクリプト /home/ssm-user/conv_bert_tensorrt.py は以下のとおりです。

#!/usr/bin/env python
# coding: utf-8

import os
import re
import pickle
import numpy as np
import torch
import torch.nn as nn
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
from transformers import AutoTokenizer, BertForTokenClassification

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model_path = "/work/hf-models/bert-ner"
max_num_labels = 45
max_length = 40
torch_script_path = "/work/trt_bert.pt"

def joint_sub_word_tokens(tokens):
  word_level_tokens = []
  for token in tokens:
    if token.startswith('##'):
      token = re.sub("^##", "", token)
      if len(word_level_tokens) > 0:
        word_level_tokens[-1]+=token
      else:
        word_level_tokens.append(token)
    else:
      word_level_tokens.append(token)
  return word_level_tokens

def build_words_to_tokens(sub_tokens):
  words_to_tokens = []
  for sub_token in sub_tokens:
    if sub_token.startswith('##'):
      words_to_tokens[-1].append(sub_token)
    else:
      words_to_tokens.append([sub_token])
  return words_to_tokens

def build_labeld_positions(words_to_tokens):
  tagged_positions = []
  pos_sub_words = 0
  for word_tokens in words_to_tokens:
    if "[CLS]" not in word_tokens and "[SEP]" not in word_tokens and "[PAD]" not in word_tokens:
      tagged_positions.append(pos_sub_words)
    pos_sub_words += len(word_tokens)
  return tagged_positions

def build_features(text, tokenizer, max_len):
  input_ids = tokenizer.encode(text, add_special_tokens=True,
                pad_to_max_length=True, max_length=max_len, truncation=True)
  sep_index = input_ids.index(tokenizer.vocab["[SEP]"])
  input_mask = [1]* len(input_ids[:sep_index+1])
  input_mask = input_mask + [0]*(max_len-len(input_mask))
  segment_ids = [0] * max_len
  sub_tokens = tokenizer.tokenize(text)
  tokens = joint_sub_word_tokens(sub_tokens)
  sub_tokens = tokenizer.convert_ids_to_tokens(input_ids)
  words_to_tokens = build_words_to_tokens(sub_tokens)
  labeled_positions =  build_labeld_positions(words_to_tokens)

  return input_ids, input_mask, segment_ids, labeled_positions, sub_tokens, tokens

def generate_input(texts, tokenizer, max_len=128):
  all_input_ids = []
  all_input_mask = []
  all_segment_ids = []
  all_labeled_positions = []
  all_sub_tokens = []
  all_tokens = []
  for i, text in enumerate(texts):
    input_ids, input_mask, segment_ids, labeled_positions, sub_tokens, tokens = build_features(text, tokenizer, max_len)
    all_input_ids.append(input_ids)
    all_input_mask.append(input_mask)
    all_segment_ids.append(segment_ids)
    # 推論時の labeled_positons の長さは未知ですが、Neuron の場合はシェイプを固定する
    # 必要があるのでシーケンスの最大長に合わせて 0 パディングしました。
    labeled_positions = labeled_positions + ([0] * (max_len - len(labeled_positions)))
    all_labeled_positions.append(labeled_positions)
    all_sub_tokens.append(sub_tokens)
    all_tokens.append(tokens)

  return (np.array(all_input_ids, dtype=np.int32),
      np.array(all_input_mask, dtype=np.int32),
      np.array(all_segment_ids, dtype=np.int32),
      np.array(all_labeled_positions, dtype=np.int32),
      np.array(all_sub_tokens),
      np.array(all_tokens))

# Convert

with torch.no_grad():

  print(f"Loading model from {model_path} to BertForTokenClassification.")
  bert = BertForTokenClassification.from_pretrained(model_path, torchscript=True)
  bert.to(device)

  class BertEncoder(nn.Module):
    def __init__(self, bert, max_num_labels):
      super(BertEncoder, self).__init__()
      self.bert = bert
      self.max_num_labels = max_num_labels

    def forward(self, input_ids, attention_mask):
      # TensorRT でコンパイルしようとすると gather() でエラーになったので、
      # flatten して末尾に max_num_labels まで 0 パディングした probs を返します。
      logits = bert.forward(input_ids, attention_mask)[0]
      probs = nn.functional.softmax(logits, dim=2)
      flatten = probs.reshape([probs.shape[0], -1])
      pad = torch.zeros([probs.shape[0], probs.shape[1]*self.max_num_labels - flatten.shape[1]], device="cuda")
      return torch.cat([flatten, pad], dim=1)

  encoder = BertEncoder(bert, max_num_labels)

  encoder.to(device).eval()

  tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)

  inputs = ["じつは東京は日本の首都です。"]

  input_ids, input_mask, segment_ids, labeled_positions, sub_tokens, tokens = generate_input(inputs,
                                                                                tokenizer, max_length)

  input_ids = torch.tensor(input_ids).to(device).to(torch.int64)
  attention_mask = torch.tensor(input_mask).to(device).to(torch.int64)
  labeled_positions = torch.tensor(labeled_positions).to(device)
  num_labels = torch.tensor([bert.config.num_labels]*len(inputs))

  import torch_tensorrt
  traced_model = torch.jit.trace(encoder, [input_ids, attention_mask])
  traced_model.to(device)
  batch_size = len(inputs)
  print(f"Invoke compile Torch-TensorRT...")
  trt_model = torch_tensorrt.compile(traced_model,
    inputs= [torch_tensorrt.Input(shape=[batch_size, max_length], dtype=torch.int32),
             torch_tensorrt.Input(shape=[batch_size, max_length], dtype=torch.int32)],
    enabled_precisions= {torch.half},
    workspace_size=2000000000,
    truncate_long_and_double=True
  )
  print(f"Saving Torch-TensorRT torch script to {torch_script_path}")
  torch.jit.save(trt_model, torch_script_path)

以下のようにして変換します。

$ cd /home/ssm-user
$ docker run --gpus all --rm -v `pwd`:/work container-d /work/conv_bert_tensorrt.py
...(大量の警告）
Loading model from /work/hf-models/bert-ner to BertForTokenClassification.
Invoke compile Torch-TensorRT...
Saving Torch-TensorRT torch script to /work/trt_bert.pt
$ ls *.pt
trt_bert.pt

大量に警告が出ましたがとりあえず変換できました。

変換スクリプト中のコメントにも書きましたが torch.gather を使うと変換エラーになったので、Triton 側ではトークン系列全体を softmax するところ( probs )まで行い、gather(), argmax() は呼び出し元で行うことにしました。

少し細かい話

probs を flatten して最後にパディングしているのは個人的な好みです。まずパディングした理由は、後述する config.pbtxt に返却するシェイプを記述する必要があることです。 config.pbtxt はモデルの全バージョンで共通なので「バージョン 3 までは 45 分類だったけど、バージョン 4 では 48 分類にしたい」というときに困ります。「分類数が違えば計算グラフも違うから別のモデルでしょ」というのも正しいのですが、利用者目線では「固有表現抽出モデルで学習したデータが違うだけ」という見方もあります。今回は後者目線で予め分類数が増えそうな分を見込んでパディングしておいて、呼び出し側で必要な分を切り出すような形にしました³¹。flatten したのはそうしないと TensorRT で変換したモデルの動きが怪しかったからだったと記憶してます。

モデルリポジトリに以下のように brt_bert.pt を model.pt にリネームして配置します。

+ trt-models/
    + trt_bert/
        + config.pbtxt
        + 1/
          + model.pt

config.pbtxt は以下のとおりです。

platform: "pytorch_libtorch"
max_batch_size: 0
input [
  {
    name: "input_ids__0"
    data_type: TYPE_INT32
    dims: [ 1, 40 ]
  },
  {
    name: "attention_mask__1"
    data_type: TYPE_INT32
    dims: [ 1, 40 ]
  }
]
output [
  {
    name: "probs__0"
    data_type: TYPE_FP32
    dims: [ 1, 180 ] # 180 = 40(seq_len) * 45(num_labels)
  }
]
version_policy: { specific { versions : [1] }}
instance_group [ {  count: 8 } ]

Torch-TensorRT の変換でバッチサイズを 1 固定としたので、max_batch_size: 0 として動的バッチングを無効にしています。その影響で入出力テンソルのシェイプの定義に [1, 40] のようにバッチ軸の値(1)が含まれているのに注意して下さい。

続いて Neuron への変換処理です。

10. Neuron への変換

作業環境は inf1.2xlarge インスタンスで transformers のモデルが /home/ssm-user/hf-models にある前提です。

Neuron は Decoder-Only, Encoder-Decoder は苦手なようなので、BERT を対象とします。

BERT の変換用スクリプト /home/ssm-user/conv_bert_neuron.py は以下のとおりです。長くなるので、Torch-TensorRT と共通の関数は省略します。

#!/usr/bin/env python
# coding: utf-8

import os
import re
import pickle
import numpy as np
import torch
import torch.nn as nn
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union
from transformers import AutoTokenizer, BertForTokenClassification

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model_path = "/work/hf-models/bert-ner"
max_num_labels = 45
max_length = 40
torch_script_path = "/work/neuron_bert.pt"

def joint_sub_word_tokens(tokens):
  # 省略

def build_words_to_tokens(sub_tokens):
  # 省略

def build_labeld_positions(words_to_tokens):
  # 省略

def build_features(text, tokenizer, max_len):
  # 省略

def generate_input(texts, tokenizer, max_len=128):
  # 省略

with torch.no_grad():

  print(f"Loading model from {model_path} to BertForTokenClassification.")
  bert = BertForTokenClassification.from_pretrained(model_path, torchscript=True)
  bert.to(device)

  class BertEncoder(nn.Module):
    def __init__(self, bert):
      super(BertEncoder, self).__init__()
      self.bert = bert

    def forward(self, input_ids, attention_mask, labeled_positions, num_labels):
      logits = bert.forward(input_ids, attention_mask)[0]
      labeled_positions = labeled_positions.unsqueeze(-1)
      index = torch.zeros(list(labeled_positions.shape[:2]) + [logits.shape[-1]]).to(device)
      index = (index + labeled_positions).to(torch.int64).to(device)
      logits = torch.gather(input=logits, dim=1, index=index)
      probs = nn.functional.softmax(logits, dim=2)
      padded = torch.zeros((probs.shape[0], probs.shape[1], num_labels[0])).to(device)
      padded[:, :, :probs.shape[2]] = probs
      probs = padded
      preds = torch.argmax(probs, axis=2)
      return preds, probs

  encoder = BertEncoder(bert)
  encoder.to(device).eval()

  tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=False)

  inputs = ["じつは東京は日本の首都です。"]

  input_ids, input_mask, segment_ids, labeled_positions, sub_tokens, tokens = generate_input(inputs,
                                                                                tokenizer, max_length)

  input_ids = torch.tensor(input_ids).to(device).to(torch.int64)
  attention_mask = torch.tensor(input_mask).to(device).to(torch.int64)
  labeled_positions = torch.tensor(labeled_positions).to(device)
  num_labels = torch.tensor([bert.config.num_labels]*len(inputs))

  import torch.neuron
  print(f"Invoke JIT tracing...")
  traced_model = torch.neuron.trace(encoder, [input_ids, attention_mask, labeled_positions, num_labels])
  print(f"Saving torch script to {torch_script_path}")
  traced_model.save(torch_script_path)
  os.system("sync")
  print("DONE")

num_labels を推論時に指定するようなコードが入っていますが、実は特に意味がないです³²。。。

最後にも怪しげな sync が入っていますが、以下のようにして変換します。

$ cd /home/ssm-user
$ docker run --rm -v `pwd`:/work container-c /work/conv_bert_neuron.py
...
NFO:Neuron:There are 4 ops of 2 different types in the TorchScript that are not compiled by neuron-cc: aten::gather, aten::embedding, (For moreinformation see https://github.com/aws/aws-neuron-sdk/blob/master/release-notes/neuron-cc-ops/neuron-cc-ops-pytorch.md)
INFO:Neuron:Number of arithmetic operators (pre-compilation) before = 598, fused = 556, percent fused = 92.98%
INFO:Neuron:Number of neuron graph operations 1631 did not match traced graph 1331 - using heuristic matching of hierarchical information
INFO:Neuron:Compiling function _NeuronGraph$689 with neuron-cc
INFO:Neuron:Compiling with command line: '/opt/conda/bin/neuron-cc compile /home/model-server/tmp/tmpzhobu3aa/graph_def.pb --framework TENSORFLOW --pipeline compile SaveTemps --output /home/model-server/tmp/tmpzhobu3aa/graph_def.neff --io-config {"inputs": {"0:0": [[1, 40, 768], "float32"], "1:0": [[1, 1, 1, 40], "float32"], "2:0": [[1, 40], "int32"]}, "outputs": ["Linear_8/aten_linear/Add:0", "aten_to_1/Cast:0"]} --verbose 35'
.....
Compiler status PASS
INFO:Neuron:skip_inference_context for tensorboard symbols at /opt/conda/lib/python3.7/site-packages/torch_neuron/tensorboard.py:305 tb_parse
INFO:Neuron:Number of neuron graph operations 1631 did not match traced graph 1331 - using heuristic matching of hierarchical information
INFO:Neuron:Number of arithmetic operators (post-compilation) before = 598, compiled = 556, percent compiled = 92.98%
INFO:Neuron:The neuron partitioner created 1 sub-graphs
INFO:Neuron:Neuron successfully compiled 1 sub-graphs, Total fused subgraphs = 1, Percent of model sub-graphs successfully compiled = 100.0%
INFO:Neuron:Compiled these operators (and operator counts) to Neuron:
INFO:Neuron: => aten::Int: 99
INFO:Neuron: => aten::add: 37
INFO:Neuron: => aten::contiguous: 12
INFO:Neuron: => aten::div: 12
INFO:Neuron: => aten::dropout: 38
INFO:Neuron: => aten::gelu: 12
INFO:Neuron: => aten::layer_norm: 25
INFO:Neuron: => aten::linear: 73
INFO:Neuron: => aten::matmul: 24
INFO:Neuron: => aten::permute: 48
INFO:Neuron: => aten::size: 99
INFO:Neuron: => aten::softmax: 12
INFO:Neuron: => aten::to: 3
INFO:Neuron: => aten::transpose: 12
INFO:Neuron: => aten::unsqueeze: 1
INFO:Neuron: => aten::view: 48
INFO:Neuron: => aten::zeros: 1
INFO:Neuron:Not compiled operators (and operator counts) to Neuron:
INFO:Neuron: => aten::Int: 8 [supported]
INFO:Neuron: => aten::add: 3 [supported]
INFO:Neuron: => aten::argmax: 1 [supported]
INFO:Neuron: => aten::copy_: 1 [supported]
INFO:Neuron: => aten::embedding: 3 [not supported]
INFO:Neuron: => aten::expand: 1 [supported]
INFO:Neuron: => aten::gather: 1 [not supported]
INFO:Neuron: => aten::mul: 1 [supported]
INFO:Neuron: => aten::rsub: 1 [supported]
INFO:Neuron: => aten::select: 1 [supported]
INFO:Neuron: => aten::size: 6 [supported]
INFO:Neuron: => aten::slice: 9 [supported]
INFO:Neuron: => aten::softmax: 1 [supported]
INFO:Neuron: => aten::to: 2 [supported]
INFO:Neuron: => aten::unsqueeze: 2 [supported]
INFO:Neuron: => aten::zeros: 1 [supported]
INFO:Neuron:skip_inference_context for tensorboard symbols at /opt/conda/lib/python3.7/site-packages/torch_neuron/tensorboard.py:305 tb_parse
INFO:Neuron:Number of neuron graph operations 139 did not match traced graph 230 - using heuristic matching of hierarchical information
Loading model from /work/hf-models/bert-ner to BertForTokenClassification.
Invoke JIT tracing...
Saving torch script to /work/neuron_bert.pt
DONE
^CTraceback (most recent call last):
  File "/usr/local/bin/dockerd-entrypoint.py", line 29, in <module>
    subprocess.call(['tail', '-f', '/dev/null'])
  File "/opt/conda/lib/python3.7/subprocess.py", line 341, in call
    return p.wait(timeout=timeout)
  File "/opt/conda/lib/python3.7/subprocess.py", line 1019, in wait
    return self._wait(timeout=timeout)
  File "/opt/conda/lib/python3.7/subprocess.py", line 1653, in _wait
    (pid, sts) = self._try_wait(0)
  File "/opt/conda/lib/python3.7/subprocess.py", line 1611, in _try_wait
    (pid, sts) = os.waitpid(self.pid, wait_flags)
KeyboardInterrupt
$ ls *.pt
neuron_bert.pt

変換には成功して保存し終わったようなのですが、いつまでたってもプロンプトに戻らないので CTRL+C で強制終了してしまいました。 traced_model.save() が終わって sync した後の “DONE” の表示があるので、多分大丈夫でしょう。本当に大丈夫かどうかは、このあと推論してみて確かめることにします。

Torch-TensorRT では gather() や argmax() が変換できなかったのですが、Neuron は綺麗に全部変換できました。 Neuron での推論には Triton を使用しなかったのでデプロイ云々の話はありません。

そうそう、前述したとおり T5 も Neuron で動かせなくはないようなので変換してみましょう。

T5 を Neuron で動かせるように変換する。

T5 を Neuron で動かすのはこちらの issue²⁴ の手順に従いました。戦略としては、

T5 の encoder と decoder をそれぞれ別個に Neuron にコンパイルして、
それらを保持する形の transformers の GenerationMixin を実装したクラス( NeuronGeneration )を用意した。

ということのようです。推論する時は NeuronGeneration としてロードし、 transformers の T5 や GPT と同様に generate() メソッドで推論できるようです。

注意

私はかなりハマったのですが、この手順は T5 1.0 のみ有効で T5 1.1 には対応していません。 なので、この連載で何度か利用した Megagon Labs さんの megagonlabs/t5-base-japanese-web は使用できません。どうしたものかと思案していたところで sonoisa さんがかなりのコーパス量で T5 1.0 のモデルを公開して下さっているのを思い出しました! そんな訳で sonoisa/t5-base-japanese をベースにファインチューニングしたモデルを hf-models/t5-1.0-base に配置した想定で記述してます。

T5 の変換用スクリプト /home/ssm-user/conv_t5_neuron.py は以下のとおりです。一部省略してますが issue²⁴ と同じ内容のコードです。このコード動作させた transformers のバージョンが不明で、とりあえず手元で動かしてみると何やらエラーが出たので NeuronGeneration は現物合わせで修正しました（__call__() で attention_mask に詰め物している箇所だったかと）。

#!/usr/bin/env python
# coding: utf-8

import os
import numpy as np
import torch
from torch.nn import functional as F
from dataclasses import dataclass, field
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union, cast
from transformers import (AutoTokenizer, HfArgumentParser, T5ForConditionalGeneration,
                          PreTrainedModel, T5Tokenizer)
from transformers.generation_utils import GenerationMixin
from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
import torch.neuron

model_path = "/work/hf-models/t5-1.0-base"
num_texts = 1
num_beams = 1
max_encoder_length=40
max_decoder_length=40
traced_model_dir="/work/neuron_t5"

with torch.no_grad():
  print(f"Loading the model from {model_path}...")
  model = cast(T5ForConditionalGeneration, T5ForConditionalGeneration.from_pretrained(model_path))
  tokenizer = cast(T5Tokenizer, T5Tokenizer.from_pretrained(model_path))

def asscaler(array):
  return array.item()
np.asscalar = asscaler

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

def reduce(hidden: torch.Tensor, index: int):
  # 省略

class NeuronEncoder(torch.nn.Module):
  # 省略

class NeuronDecoder(torch.nn.Module):
  # 省略

class NeuronGeneration(PreTrainedModel, GenerationMixin):
  def trace(self, model: T5ForConditionalGeneration, num_texts: int, num_beams: int,
        max_encoder_length: int, max_decoder_length: int) -> None:
    self.config.max_decoder_length = max_decoder_length
    inputs = (
      torch.ones((num_texts, max_encoder_length), dtype=torch.long),
      torch.ones((num_texts, max_encoder_length), dtype=torch.long),
    )
    encoder = NeuronEncoder(model)
    self.encoder = cast(NeuronEncoder, torch.neuron.trace(encoder, inputs))
    batch_size = num_texts * num_beams
    inputs = (
      torch.ones((batch_size, max_decoder_length), dtype=torch.long),
      torch.ones((batch_size, max_encoder_length), dtype=torch.long),
      torch.ones( (batch_size, max_encoder_length, model.config.d_model), dtype=torch.float,),
      torch.tensor(0),
    )
    decoder = NeuronDecoder(model, max_decoder_length)
    self.decoder = cast(NeuronDecoder, torch.neuron.trace(decoder, inputs))

  def prepare_inputs_for_generation(self, input_ids: torch.Tensor, encoder_outputs: BaseModelOutput,
        attention_mask: Optional[BaseModelOutput] = None, **model_kwargs: Any):
    current_length = input_ids.shape[1]
    pad_size = self.config.max_decoder_length - current_length
    return dict(
      input_ids=F.pad(input_ids, (0, pad_size)),
      attention_mask=attention_mask,
      encoder_outputs=encoder_outputs.last_hidden_state,
      current_length=torch.tensor(current_length - 1),
    )

  def get_encoder(self):
    def encode(**kwargs:Any):
      input_ids = kwargs["input_ids"]
      attention_mask = kwargs.get("attention_mask", torch.ones_like(input_ids))
      (output,) = self.encoder(input_ids, attention_mask)
      return BaseModelOutput(last_hidden_state=output)
    return encode

  def __call__(self, input_ids: torch.Tensor, attention_mask: torch.Tensor,
        encoder_outputs: BaseModelOutput, current_length: int, **kwargs: Any):
    if attention_mask is None:
      # attention_mask will be ignored in NeuronDecoder, the value is not a matter here.
      attention_mask = torch.zeros_like(input_ids)
    logits = self.decoder(input_ids, attention_mask, encoder_outputs, current_length)
    return Seq2SeqLMOutput(logits=logits)

  def save_pretrained(self, directory: str):
    if os.path.isfile(directory):
      print(f"Provided path ({directory}) should be a directory, not a file")
      return
    os.makedirs(directory, exist_ok=True)
    torch.jit.save(self.encoder, os.path.join(directory, "encoder.pt"))
    torch.jit.save(self.decoder, os.path.join(directory, "decoder.pt"))
    self.config.save_pretrained(directory)

  @classmethod
  def from_pretrained(cls, directory: str):
    config = T5Config.from_pretrained(directory)
    obj = cls(config)
    obj.main_input_name = "input_ids"
    obj.encoder = torch.jit.load(os.path.join(directory, "encoder.pt"))
    obj.encoder.main_input_name = "input_ids"
    obj.decoder = torch.jit.load(os.path.join(directory, "decoder.pt"))
    obj.decoder.main_input_name = "decoder_input_ids"
    return obj

  @property
  def device(self):
    return torch.device("cpu")

with torch.no_grad():
  print(f"Loading the model from {model_path}...")
  model = cast(T5ForConditionalGeneration, T5ForConditionalGeneration.from_pretrained(model_path))
  tokenizer = cast(T5Tokenizer, T5Tokenizer.from_pretrained(model_path))
  model_neuron = NeuronGeneration(model.config)

  print("Start tracing...")
  model_neuron.trace(model=model,
    num_texts=num_texts,
    num_beams=num_beams,
    max_encoder_length=max_encoder_length,
    max_decoder_length=max_decoder_length,)

  print(f"Saving traced model to {traced_model_dir}...")
  model_neuron.save_pretrained(traced_model_dir)
  os.system("sync")
  print("DONE")

以下のようにして変換します。

$ docker run --rm -v `pwd`:/work container-c /work/conv_t5_neuron.py
INFO:Neuron:There are 2 ops of 1 different types in the TorchScript that are not compiled by neuron-cc: aten::embedding, (For more information see https://github.com/aws/aws-neuron-sdk/blob/master/release-notes/neuron-cc-ops/neuron-cc-ops-pytorch.md)
INFO:Neuron:Number of arithmetic operators (pre-compilation) before = 627, fused = 587, percent fused = 93.62%
INFO:Neuron:Number of neuron graph operations 1637 did not match traced graph 1569 - using heuristic matching of hierarchical information
WARNING:tensorflow:From /opt/conda/lib/python3.7/site-packages/torch_neuron/ops/aten.py:2022: where (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.
Instructions for updating:
Use tf.where in 2.0, which has the same broadcast rule as np.where
INFO:Neuron:Compiling function _NeuronGraph$641 with neuron-cc
INFO:Neuron:Compiling with command line: '/opt/conda/bin/neuron-cc compile /home/model-server/tmp/tmpb7eevuw1/graph_def.pb --framework TENSORFLOW --pipeline compile SaveTemps --output /home/model-server/tmp/tmpb7eevuw1/graph_def.neff --io-config {"inputs": {"tensor.1:0": [[1, 40], "int64"], "1:0": [[1, 40, 768], "float32"], "2:0": [[40, 40, 12], "float32"], "3:0": [[1, 40, 768], "float32"]}, "outputs": ["T5Stack_1/T5LayerNorm_79/aten_mul_1/mul:0"]} --verbose 35'
....
Compiler status PASS
INFO:Neuron:skip_inference_context for tensorboard symbols at /opt/conda/lib/python3.7/site-packages/torch_neuron/tensorboard.py:305 tb_parse
INFO:Neuron:Number of neuron graph operations 1637 did not match traced graph 1569 - using heuristic matching of hierarchical information
INFO:Neuron:Number of arithmetic operators (post-compilation) before = 627, compiled = 587, percent compiled = 93.62%
INFO:Neuron:The neuron partitioner created 1 sub-graphs
INFO:Neuron:Neuron successfully compiled 1 sub-graphs, Total fused subgraphs = 1, Percent of model sub-graphs successfully compiled = 100.0%
INFO:Neuron:Compiled these operators (and operator counts) to Neuron:
INFO:Neuron: => aten::Int: 48
INFO:Neuron: => aten::add: 61
INFO:Neuron: => aten::contiguous: 12
INFO:Neuron: => aten::dropout: 49
INFO:Neuron: => aten::linear: 72
INFO:Neuron: => aten::matmul: 24
INFO:Neuron: => aten::mean: 24
INFO:Neuron: => aten::mul: 49
INFO:Neuron: => aten::permute: 1
INFO:Neuron: => aten::pow: 24
INFO:Neuron: => aten::relu: 12
INFO:Neuron: => aten::rsqrt: 24
INFO:Neuron: => aten::rsub: 1
INFO:Neuron: => aten::size: 12
INFO:Neuron: => aten::slice: 2
INFO:Neuron: => aten::softmax: 12
INFO:Neuron: => aten::to: 37
INFO:Neuron: => aten::transpose: 60
INFO:Neuron: => aten::type_as: 12
INFO:Neuron: => aten::unsqueeze: 3
INFO:Neuron: => aten::view: 48
INFO:Neuron:Not compiled operators (and operator counts) to Neuron:
INFO:Neuron: => aten::Int: 1 [supported]
INFO:Neuron: => aten::ScalarImplicit: 2 [supported]
INFO:Neuron: => aten::abs: 1 [supported]
INFO:Neuron: => aten::add: 4 [supported]
INFO:Neuron: => aten::arange: 2 [supported]
INFO:Neuron: => aten::div: 2 [supported]
INFO:Neuron: => aten::dropout: 1 [supported]
INFO:Neuron: => aten::embedding: 2 [not supported]
INFO:Neuron: => aten::full_like: 1 [supported]
INFO:Neuron: => aten::gt: 1 [supported]
INFO:Neuron: => aten::log: 1 [supported]
INFO:Neuron: => aten::lt: 1 [supported]
INFO:Neuron: => aten::mean: 1 [supported]
INFO:Neuron: => aten::min: 1 [supported]
INFO:Neuron: => aten::mul: 4 [supported]
INFO:Neuron: => aten::pow: 1 [supported]
INFO:Neuron: => aten::rsqrt: 1 [supported]
INFO:Neuron: => aten::size: 2 [supported]
INFO:Neuron: => aten::slice: 2 [supported]
INFO:Neuron: => aten::sub: 1 [supported]
INFO:Neuron: => aten::to: 4 [supported]
INFO:Neuron: => aten::unsqueeze: 2 [supported]
INFO:Neuron: => aten::view: 1 [supported]
INFO:Neuron: => aten::where: 1 [not supported]
INFO:Neuron:skip_inference_context for tensorboard symbols at /opt/conda/lib/python3.7/site-packages/torch_neuron/tensorboard.py:305 tb_parse
INFO:Neuron:Number of neuron graph operations 119 did not match traced graph 120 - using heuristic matching of hierarchical information
/opt/conda/lib/python3.7/site-packages/transformers/modeling_utils.py:781: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if causal_mask.shape[1] < attention_mask.shape[1]:
INFO:Neuron:There are 3 ops of 2 different types in the TorchScript that are not compiled by neuron-cc: aten::embedding, aten::repeat, (For moreinformation see https://github.com/aws/aws-neuron-sdk/blob/master/release-notes/neuron-cc-ops/neuron-cc-ops-pytorch.md)
INFO:Neuron:Number of arithmetic operators (pre-compilation) before = 1094, fused = 1035, percent fused = 94.61%
INFO:Neuron:Number of neuron graph operations 2895 did not match traced graph 2814 - using heuristic matching of hierarchical information
INFO:Neuron:Compiling function _NeuronGraph$1572 with neuron-cc
INFO:Neuron:Compiling with command line: '/opt/conda/bin/neuron-cc compile /home/model-server/tmp/tmpcauw6xsv/graph_def.pb --framework TENSORFLOW --pipeline compile SaveTemps --output /home/model-server/tmp/tmpcauw6xsv/graph_def.neff --io-config {"inputs": {"tensor.1:0": [[1, 40, 40], "bool"], "tensor.9:0": [[1, 40], "float32"], "tensor.15:0": [[1, 40], "int64"], "3:0": [[1, 40, 768], "float32"], "4:0": [[40, 40, 12], "float32"], "5:0": [[1, 40, 768], "float32"], "6:0": [[1, 40, 768], "float32"], "7:0": [[], "int64"]}, "outputs": ["aten_linear/MatMul:0"]} --verbose 35'
.......
Compiler status PASS
INFO:Neuron:skip_inference_context for tensorboard symbols at /opt/conda/lib/python3.7/site-packages/torch_neuron/tensorboard.py:305 tb_parse
INFO:Neuron:Number of neuron graph operations 2895 did not match traced graph 2814 - using heuristic matching of hierarchical information
INFO:Neuron:Number of arithmetic operators (post-compilation) before = 1094, compiled = 1035, percent compiled = 94.61%
INFO:Neuron:The neuron partitioner created 1 sub-graphs
INFO:Neuron:Neuron successfully compiled 1 sub-graphs, Total fused subgraphs = 1, Percent of model sub-graphs successfully compiled = 100.0%
INFO:Neuron:Compiled these operators (and operator counts) to Neuron:
INFO:Neuron: => aten::Int: 98
INFO:Neuron: => aten::ScalarImplicit: 1
INFO:Neuron: => aten::add: 98
INFO:Neuron: => aten::arange: 1
INFO:Neuron: => aten::contiguous: 24
INFO:Neuron: => aten::dropout: 73
INFO:Neuron: => aten::eq: 1
INFO:Neuron: => aten::linear: 121
INFO:Neuron: => aten::matmul: 48
INFO:Neuron: => aten::mean: 36
INFO:Neuron: => aten::mul: 76
INFO:Neuron: => aten::permute: 1
INFO:Neuron: => aten::pow: 36
INFO:Neuron: => aten::relu: 12
INFO:Neuron: => aten::rsqrt: 36
INFO:Neuron: => aten::rsub: 2
INFO:Neuron: => aten::size: 27
INFO:Neuron: => aten::slice: 7
INFO:Neuron: => aten::softmax: 24
INFO:Neuron: => aten::sum: 1
INFO:Neuron: => aten::to: 63
INFO:Neuron: => aten::transpose: 120
INFO:Neuron: => aten::type_as: 24
INFO:Neuron: => aten::unsqueeze: 7
INFO:Neuron: => aten::view: 97
INFO:Neuron: => aten::zeros: 1
INFO:Neuron:Not compiled operators (and operator counts) to Neuron:
INFO:Neuron: => aten::Int: 7 [supported]
INFO:Neuron: => aten::ScalarImplicit: 3 [supported]
INFO:Neuron: => aten::add: 3 [supported]
INFO:Neuron: => aten::arange: 3 [supported]
INFO:Neuron: => aten::div: 2 [supported]
INFO:Neuron: => aten::dropout: 1 [supported]
INFO:Neuron: => aten::embedding: 2 [not supported]
INFO:Neuron: => aten::full_like: 1 [supported]
INFO:Neuron: => aten::le: 1 [supported]
INFO:Neuron: => aten::log: 1 [supported]
INFO:Neuron: => aten::lt: 1 [supported]
INFO:Neuron: => aten::mean: 1 [supported]
INFO:Neuron: => aten::min: 2 [supported]
INFO:Neuron: => aten::mul: 3 [supported]
INFO:Neuron: => aten::neg: 1 [supported]
INFO:Neuron: => aten::ones: 2 [supported]
INFO:Neuron: => aten::pow: 1 [supported]
INFO:Neuron: => aten::repeat: 1 [not supported]
INFO:Neuron: => aten::rsqrt: 1 [supported]
INFO:Neuron: => aten::size: 4 [supported]
INFO:Neuron: => aten::slice: 4 [supported]
INFO:Neuron: => aten::sub: 1 [supported]
INFO:Neuron: => aten::to: 4 [supported]
INFO:Neuron: => aten::unsqueeze: 6 [supported]
INFO:Neuron: => aten::view: 1 [supported]
INFO:Neuron: => aten::where: 1 [not supported]
INFO:Neuron: => aten::zeros_like: 1 [supported]
INFO:Neuron:skip_inference_context for tensorboard symbols at /opt/conda/lib/python3.7/site-packages/torch_neuron/tensorboard.py:305 tb_parse
INFO:Neuron:Number of neuron graph operations 175 did not match traced graph 182 - using heuristic matching of hierarchical information
Loading the model from /work/hf-models/t5-1.0-base...
Loading the model from /work/hf-models/t5-1.0-base...
Start tracing...
Saving traced model to /work/neuron_t5...
DONE
^CTraceback (most recent call last):
  File "/usr/local/bin/dockerd-entrypoint.py", line 29, in <module>
    subprocess.call(['tail', '-f', '/dev/null'])
  File "/opt/conda/lib/python3.7/subprocess.py", line 341, in call
    return p.wait(timeout=timeout)
  File "/opt/conda/lib/python3.7/subprocess.py", line 1019, in wait
    return self._wait(timeout=timeout)
  File "/opt/conda/lib/python3.7/subprocess.py", line 1653, in _wait
    (pid, sts) = self._try_wait(0)
  File "/opt/conda/lib/python3.7/subprocess.py", line 1611, in _try_wait
    (pid, sts) = os.waitpid(self.pid, wait_flags)
KeyboardInterrupt
$ ls neuron_t5/
config.json  decoder.pt  encoder.pt

こちらも CTRL+C で止めましたが、どうやら変換出来ているみたいです。

これでモデルの変換は完了です。

調子にのっていろいろ作りこんだら思いのほか gunicorn / Falcon のコードが長くなって、速度計測に使ったコードは全部載せられないのですが、参考までに各モデルの推論コードをご紹介します。

11. FasterTransformer での推論

まずは FasterTransformer に変換したモデルで推論をしてみましょう。

先程の g4dn.xlarge インスタンスで /home/ssm-user/models が FasterTransformer に変換したモデルのモデルリポジトリだとします。

$ pwd 
/home/ssm-user
$ ls
fastertransformer_backend  hf-models models

models の内容物はこんな感じです。

+ models/
    + ft_bert/
        + config.pbtxt
        + 1/
          + 1-gpu/
    + ft_gpt2/
        + config.pbtxt
        + 1/
          + 1-gpu/
    + ft_t5/
        + config.pbtxt
        + 1/
          + 1-gpu/

まずは Triton を起動します。

$ docker run --name triton -d --rm --gpus all \
    --ulimit memlock=-1 --ulimit stack=67108864 --shm-size=1g \
    -v /home/ssm-user/models:/models \
    container-a2 tritonserver --model-repository=/models
6e887b5c9e655c01b93e8d9c8f373447e4373c04468849d8cc5bd301cf087436
$ docker logs triton
=============================
== Triton Inference Server ==
=============================

NVIDIA Release 22.07 (build 41737377)
Triton Server Version 2.24.0
...
I0123 04:23:57.463060 1 server.cc:559]
+------------------+------+
| Repository Agent | Path |
+------------------+------+
+------------------+------+

I0123 04:23:57.463115 1 server.cc:586]
+-------------------+-----------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------+
| Backend           | Path                                                                        | Config                                                                                                                 |
+-------------------+-----------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------+
| fastertransformer | /opt/tritonserver/backends/fastertransformer/libtriton_fastertransformer.so | {"cmdline":{"auto-complete-config":"true","min-compute-capability":"6.000000","backend-directory":"/opt/tritonserver/backends","default-max-batch-size":"4"}} |
+-------------------+-----------------------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------+

I0123 04:23:57.463156 1 server.cc:629]
+---------+---------+--------+
| Model   | Version | Status |
+---------+---------+--------+
| ft_bert | 1       | READY  |
| ft_gpt2 | 1       | READY  |
| ft_t5   | 1       | READY  |
+---------+---------+--------+

I0123 04:23:57.478870 1 metrics.cc:650] Collecting metrics for GPU 0: Tesla T4
I0123 04:23:57.479103 1 tritonserver.cc:2176]
+----------------------------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| Option                           | Value                                                                                 |
+----------------------------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| server_id                        | triton                                                                                 |
| server_version                   | 2.24.0                                                                                 |
| server_extensions                | classification sequence model_repository model_repository(unload_dependents) schedule_policy model_configuration system_shared_memory cuda_shared_memory binary_tensor_data statistics trace |
| model_repository_path[0]         | /models                                                                                 |
| model_control_mode               | MODE_NONE                                                                                 |
| strict_model_config              | 0                                                                                 |
| rate_limit                       | OFF                                                                                 |
| pinned_memory_pool_byte_size     | 268435456                                                                                 |
| cuda_memory_pool_byte_size{0}    | 67108864                                                                                 |
| response_cache_byte_size         | 0                                                                                 |
| min_supported_compute_capability | 6.0                                                                                 |
| strict_readiness                 | 1                                                                                 |
| exit_timeout                     | 30                                                                                 |
+----------------------------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+

I0123 04:23:57.479908 1 grpc_server.cc:4608] Started GRPCInferenceService at 0.0.0.0:8001
I0123 04:23:57.480151 1 http_server.cc:3312] Started HTTPService at 0.0.0.0:8000
I0123 04:23:57.521108 1 http_server.cc:178] Started Metrics Service at 0.0.0.0:8002
$

BERT での推論

推論スクリプトの invoke_ft_bert.py は以下のようになりました。一部省略した箇所は本記事あるいはコメントのリンク先の同名関数と同じ内容です。

#!/usr/bin/env python
# coding: utf-8

import os
import pprint
import re
import numpy as np
import torch
import torch.nn as nn
from dataclasses import dataclass, field
from typing import Optional
from transformers import AutoTokenizer, HfArgumentParser, BertForTokenClassification
import tritonclient.grpc as grpcclient
from tritonclient.utils import np_to_triton_dtype

pp = pprint.PrettyPrinter(indent=2)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

@dataclass
class Arguments:
    host : Optional[str] = field(default='triton:8001')
    model_path : Optional[str] = field(default='/work/hf-models/bert-ner')

parser = HfArgumentParser((Arguments))
args = parser.parse_args_into_dataclasses()[0]

triton_client = grpcclient.InferenceServerClient(url=args.host, verbose=False)

def joint_sub_word_tokens(tokens):
  # 省略

def build_words_to_tokens(sub_tokens):
  # 省略

def build_labeld_positions(words_to_tokens):
  # 省略

def build_features(text, tokenizer, max_len):
  # 省略

def generate_input(texts, tokenizer, max_len=128):
  # 省略

def get_span_labels(sentence_tags, inv_label_mapping=None):
  # https://github.com/google-research/electra/blob/79111328070e491b287c307906701ebc61091eb2/finetune/tagging/tagging_utils.py#L23-L40

def get_entities(span_labels, tokens, probs):
  # https://www.ogis-ri.co.jp/otc/hiroba/technical/similar-document-search/part12.html

def build_result(all_tokens, all_sub_tokens, all_predicts, all_probs, id2label):
  results = []
  for tokens, predicts, probs in zip(all_tokens, all_predicts, all_probs):
    span_labels = get_span_labels(predicts, id2label)
    entities = get_entities(span_labels, tokens, probs)
    results.append(entities)
  return results

def prepare_tensor(name, input, protocol):
  client_util = httpclient if protocol == "http" else grpcclient
  t = client_util.InferInput(
      name, input.shape, np_to_triton_dtype(input.dtype))
  t.set_data_from_numpy(input)
  return t

def send_request(input_hidden_state, sequence_lengths, protocol="grpc"):
  inputs = [
    prepare_tensor("input_hidden_state", input_hidden_state, protocol),
    prepare_tensor("sequence_lengths", sequence_lengths, protocol),
  ]
  model_name = "ft_bert"
  result = triton_client.infer(model_name, inputs)
  output_hidden_state = result.as_numpy("output_hidden_state")
  return output_hidden_state

def main():
  print(f"Loading model from {args.model_path} to BertForTokenClassification.")
  model = BertForTokenClassification.from_pretrained(args.model_path, torchscript=True)
  model.to(device).eval()

  print(f"Loading tokenizer from {args.model_path} to BertForTokenClassification.")
  tokenizer = AutoTokenizer.from_pretrained(args.model_path)

  inputs = ["ジョー・バイデンはアメリカの大統領です。"]

  print(f"input texts : {inputs}")
  print(f"Building input features.")
  input_ids, input_mask, segment_ids, labeled_positions, sub_tokens, tokens  = generate_input(inputs, tokenizer)

  batch_size = len(input_ids)
  inputs = []
  outputs = []

  print(f"Generating input hidden states.")
  input_ids = torch.tensor(input_ids).to(device)
  token_type_ids = torch.tensor(segment_ids).to(device)
  input_hidden_state = model.bert.embeddings(input_ids=input_ids, token_type_ids=token_type_ids)
  input_hidden_state = input_hidden_state.cpu().detach().numpy().astype(np.float16)
  sequence_lengths = np.sum(input_mask, axis=1, keepdims=True).astype(np.int32)

  print(f"Invoking fastertransformer BERT.")
  ft_output_hidden_state = send_request(input_hidden_state, sequence_lengths)

  output_hidden_state = torch.tensor(ft_output_hidden_state.astype(np.float32))
  output_hidden_state = output_hidden_state.to(device)

  print(f"Invoking classifier and softmax.")
  labeled_positions = torch.tensor(labeled_positions).to(torch.int64)
  labeled_positions = labeled_positions.unsqueeze(-1)
  index = torch.zeros(list(labeled_positions.shape[:2]) + [output_hidden_state.shape[-1]])
  index = (index + labeled_positions).to(torch.int64)
  index = index.to(device)

  output_hidden_state = torch.gather(input=output_hidden_state, dim=1, index=index)
  logits = model.classifier(output_hidden_state)
  probs = nn.functional.softmax(logits, dim=2)
  preds = torch.argmax(probs, axis=2)

  probs = probs.cpu().detach().numpy()
  preds = preds.cpu().detach().numpy()

  print(f"Building final results.")
  results = build_result(tokens, sub_tokens, preds, probs, model.config.id2label)
  pp.pprint(results)

if __name__ == "__main__":
  with torch.no_grad():
    main()

BertForTokenClassification の BertModel 相当の部分だけ FasterTransformer に置き換える形なので、どうしてもスッキリしない感じになってしまいますね。。。

以下のようにして呼び出します。

$ docker run --gpus all --rm --link triton -v `pwd`:/work container-a1 /work/invoke_ft_bert.py
...
Loading model from /work/hf-models/bert-ner to BertForTokenClassification.
Loading tokenizer from /work/hf-models/bert-ner to BertForTokenClassification.
input texts : ['ジョー・バイデンはアメリカの大統領です。']
Building input features.
Generating input hidden states.
Invoking fastertransformer BERT.
Invoking classifier and softmax.
Building final results.
[ [ {'end': 8, 'start': 0, 'text': 'ジョー・バイデン', 'type': 'PERSON'},
    {'end': 13, 'start': 9, 'text': 'アメリカ', 'type': 'GPE'}]]

T5 での推論

推論スクリプト invoke_ft_t5.py は以下のとおりです。

#!/usr/bin/env python
# coding: utf-8

import numpy as np
import torch
import datetime
from dataclasses import dataclass, field
from typing import Optional
from transformers import T5Tokenizer, HfArgumentParser
import tritonclient.grpc as grpcclient
from tritonclient.utils import np_to_triton_dtype

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

@dataclass
class Arguments:
    host : Optional[str] = field(default='triton:8001')
    model_path : Optional[str] = field(default='/work/hf-models/t5-base')

parser = HfArgumentParser((Arguments))
args = parser.parse_args_into_dataclasses()[0]

triton_client = grpcclient.InferenceServerClient(url=args.host, verbose=False)

def prepare_tensor(name, input, protocol):
  client_util = httpclient if protocol == "http" else grpcclient
  t = client_util.InferInput(
      name, input.shape, np_to_triton_dtype(input.dtype))
  t.set_data_from_numpy(input)
  return t

def tokenize(tokenizer, text, max_len):
  return tokenizer(text, return_tensors='np', truncation=True, max_length=max_len,
            padding=True, pad_to_multiple_of=8)

def build_inputs(input_ids, mem_seq_len, max_output_len,
               runtime_top_k, runtime_top_p, beam_search_diversity_rate,
               temperature, len_penalty, repetition_penalty, random_seed,
               is_return_log_probs, beam_width, start_ids, end_ids,
               bad_words_ids, stop_words_ids, protocol="grpc"):
  inputs = [
      prepare_tensor("input_ids", input_ids, protocol),
      prepare_tensor("sequence_length", mem_seq_len, protocol),
      prepare_tensor("max_output_len", max_output_len, protocol),
      prepare_tensor("runtime_top_k", runtime_top_k, protocol),
      prepare_tensor("runtime_top_p", runtime_top_p, protocol),
      prepare_tensor("beam_search_diversity_rate", beam_search_diversity_rate, protocol),
      prepare_tensor("temperature", temperature, protocol),
      prepare_tensor("len_penalty", len_penalty, protocol),
      prepare_tensor("repetition_penalty", repetition_penalty, protocol),
      prepare_tensor("random_seed", random_seed, protocol),
      prepare_tensor("is_return_log_probs", is_return_log_probs, protocol),
      prepare_tensor("beam_width", beam_width, protocol),
      prepare_tensor("start_id", start_ids, protocol),
      prepare_tensor("end_id", end_ids, protocol),
      prepare_tensor("bad_words_list", bad_words_ids, protocol),
      prepare_tensor("stop_words_list", stop_words_ids, protocol),
  ]
  return inputs

def decode(tokenizer, text, result):
  output = result.as_numpy("output_ids")
  ft_output_len = result.as_numpy("sequence_length")
  return tokenizer.decode(output[0][0][:ft_output_len[0][0]], skip_special_tokens=True)

def main(model_name, tokenize, build_inputs, text):
  print(f"input text: {text}")
  print(f"Loading tokenizer from {args.model_path}")
  tokenizer = T5Tokenizer.from_pretrained(args.model_path)
  bos_token_id = tokenizer.bos_token_id if tokenizer.bos_token_id else 0
  eos_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id else 1
  print(f"bos_token_id = {bos_token_id}")
  print(f"eos_token_id = {eos_token_id}")

  print(f"Tokenizing input.")
  input_tokens = tokenize(tokenizer, text, max_len=40)

  input_ids = input_tokens.input_ids.astype(np.uint32)
  mem_seq_len = np.sum(input_tokens.attention_mask, axis=1).astype(np.uint32)
  mem_seq_len = mem_seq_len.reshape([mem_seq_len.shape[0], 1])

  runtime_top_k = (1 * np.ones([input_ids.shape[0], 1])).astype(np.uint32)
  runtime_top_p = 1.0 * np.ones([input_ids.shape[0], 1]).astype(np.float32)
  beam_search_diversity_rate = 0.0 *  np.ones([input_ids.shape[0], 1]).astype(np.float32)
  temperature = 1.0 * np.ones([input_ids.shape[0], 1]).astype(np.float32)
  len_penalty = 1.0 * np.ones([input_ids.shape[0], 1]).astype(np.float32)
  repetition_penalty = 1.0 * np.ones([input_ids.shape[0], 1]).astype(np.float32)
  random_seed = int(datetime.datetime.now().timestamp()*100)
  random_seed = random_seed * np.ones([input_ids.shape[0], 1]).astype(np.uint64)
  is_return_log_probs = False * np.ones([input_ids.shape[0], 1]).astype(bool)
  max_output_len = (64 * np.ones([input_ids.shape[0], 1])).astype(np.uint32)
  bad_words_ids = np.array([[[0], [-1]]] * input_ids.shape[0], dtype=np.int32)
  stop_words_ids = np.array([[[0], [-1]]] * input_ids.shape[0], dtype=np.int32)
  beam_width = (1 * np.ones([input_ids.shape[0], 1])).astype(np.uint32)
  start_ids = bos_token_id * np.ones([input_ids.shape[0], 1]).astype(np.uint32)
  end_ids = eos_token_id * np.ones([input_ids.shape[0], 1]).astype(np.uint32)

  print(f"Building input data.")
  inputs = build_inputs(input_ids, mem_seq_len, max_output_len,
               runtime_top_k, runtime_top_p, beam_search_diversity_rate,
               temperature, len_penalty, repetition_penalty, random_seed,
               is_return_log_probs, beam_width, start_ids, end_ids,
               bad_words_ids, stop_words_ids)

  print(f"Invoking Fastartransformer {model_name}.")
  result = triton_client.infer(model_name, inputs)

  print(f"Decoding output.")
  print(decode(tokenizer, text, result))

if __name__ == "__main__":
  with torch.no_grad():
    main("ft_t5", tokenize, build_inputs, "その男は窃盗犯に問われた。")

以下のようにして推論します。

$ docker run --rm --link triton -v `pwd`:/work container-a1 /work/invoke_ft_t5.py
...
input text: その男は窃盗犯に問われた。
Loading tokenizer from /work/hf-models/t5-base
bos_token_id = 0
eos_token_id = 1
Tokenizing input.
Building input data.
Invoking Fastartransformer ft_t5.
Decoding output.
その男は盗んだ犯罪者に問われた。

モデルは本連載で何度か行っているやさしい日本語変換です。ちょっと微妙ですが変換できてますね。

GPT2 での推論

GPT2 の推論は T5 とあまり変わりません。推論スクリプト invoke_ft_gpt2.py は以下のとおりです。

#!/usr/bin/env python
# coding: utf-8

import torch
import sys
sys.path.append(".")
from invoke_ft_t5 import prepare_tensor, main

def tokenize(tokenizer, text, max_len):
  return tokenizer(text, return_tensors='np', add_special_tokens=False, truncation=True,
                     max_length=max_len, padding=True, pad_to_multiple_of=8)

def build_inputs(input_ids, mem_seq_len, max_output_len,
               runtime_top_k, runtime_top_p, beam_search_diversity_rate,
               temperature, len_penalty, repetition_penalty, random_seed,
               is_return_log_probs, beam_width, start_ids, end_ids,
               bad_words_ids, stop_words_ids, protocol="grpc"):
  inputs = [
      prepare_tensor("input_ids", input_ids, protocol),
      prepare_tensor("input_lengths", mem_seq_len, protocol),
      prepare_tensor("request_output_len", max_output_len, protocol),
      prepare_tensor("runtime_top_k", runtime_top_k, protocol),
      prepare_tensor("runtime_top_p", runtime_top_p, protocol),
      prepare_tensor("beam_search_diversity_rate", beam_search_diversity_rate, protocol),
      prepare_tensor("temperature", temperature, protocol),
      prepare_tensor("len_penalty", len_penalty, protocol),
      prepare_tensor("repetition_penalty", repetition_penalty, protocol),
      prepare_tensor("random_seed", random_seed, protocol),
      prepare_tensor("is_return_log_probs", is_return_log_probs, protocol),
      prepare_tensor("beam_width", beam_width, protocol),
      prepare_tensor("start_id", start_ids, protocol),
      prepare_tensor("end_id", end_ids, protocol),
      prepare_tensor("bad_words_list", bad_words_ids, protocol),
      prepare_tensor("stop_words_list", stop_words_ids, protocol),
  ]
  return inputs

if __name__ == "__main__":
  with torch.no_grad():
    main("ft_gpt2", tokenize, build_inputs, "彼は寛大であると聞いている。::")

推論してみます。

$ docker run --rm --link triton -v `pwd`:/work container-a1 /work/invoke_ft_gpt2.py \
    --model_path /work/hf-models/gpt2-medium
...
input text: 彼は寛大であると聞いている。::
Loading tokenizer from /work/hf-models/gpt2-medium
bos_token_id = 1
eos_token_id = 2
Tokenizing input.
Building input data.
Invoking Fastartransformer ft_gpt2.
Decoding output.
彼は寛大であると聞いている。:: 彼は心が広いと聞いている。

このモデルは学習時に “変換元テキスト::変換先テキスト” のフォーマットで学習して Seq2Seq 風の変換が出来るようにファインチューニングしました。いい感じの変換になっています。

補足 : bad/stop_words_list

T5, GPT2 の推論で分かりにくいのは bad/stop_words_list になるかと思います。前者は指定した単語の生成を抑制する設定、後者は指定した単語が生成された段階でテキスト生成を終了する設定です。

具体例で示した方が分かりやすいですね。例えば、あなたが小動物が大嫌いだとします。「子犬」とか「子猫」なんて見たくもありません（どんなキャラ設定なんだか）。まずトークン ID を調べます。

from transformers import T5Tokenizer
tokenizer = T5Tokenizer.from_pretrained("rinna/japanese-gpt2-medium")
input_ids = tokenizer("私は子犬と子猫が好きです。").input_ids
tokens = tokenizer.convert_ids_to_tokens(input_ids)
for i, t in zip(input_ids, tokens):
  print(f"{t}=>{i}")
# ▁=>9
# 私は=>6057
# 子=>129
# 犬=>1651
# と=>20
# 子=>129
# 猫=>4324
# が好き=>13851
# です=>2767
# 。=>8
# </s>=>2

子犬は [129, 1651]、子猫は [129, 4324] です。この場合 bad_words_list に設定するテンソルは以下のようになります。

bad_words_list
# array([[[ 129, 1651,  129, 4324],
#         [   2,    4,   -1,   -1]]], dtype=int32)

テンソルのシェイプは [1, 2, 4] で一番外側はバッチサイズです。バッチの各要素に対応する中身は以下のとおりです。

bad_word_list[0][0] : bad/stop_words に指定したい単語を構成するトークンIDを flatten した物になります。
bad_word_list[0][1] : bad/stop_words の各単語を構成するトークン数の累積和です。「子犬」、「子猫」が共に 2 トークンなので [2, 4] となり余りは -1 でパディングして長さを揃えます。

次は TensorRT で BERT の推論です。

12. Torch-TensorRT での推論

FasterTransformer を動かしたのと同じ g4dn.xlarge インスタンスで /home/ssm-user/trt-models が Torch-TensorRT に変換したモデルのモデルリポジトリだとします。

今、こんな状況です（大分散らかってきました）。

$ pwd
/home/ssm-user
$ ls
build                  fastertransformer_backend  invoke_ft_bert.py  invoke_ft_t5.py  __pycache__
conv_bert_tensorrt.py  hf-models                  invoke_ft_gpt2.py  models           trt_models

trt-models の内容物はこんな感じです。model.pt は Tortch-TensorRT への変換で生成した trt_bert.pt をリネームしたものです。

+ trt-models/
    + trt_bert/
        + config.pbtxt
        + 1/
          + model.pt

まずは Triton を起動します。Torch-TensorRT にしたモデルは複数インスタンスデプロイして並行で推論できるので、景気よく 8 個デプロイしました。

$ docker run --name triton -d --rm --gpus all --ulimit memlock=-1 \
  --ulimit stack=67108864 --shm-size=1g \
  -v /home/ssm-user/trt-models:/models \
  container-a2 tritonserver --model-repository=/models
6e887b5c9e655c01b93e8d9c8f373447e4373c04468849d8cc5bd301cf087436
$ docker logs triton

=============================
== Triton Inference Server ==
=============================

NVIDIA Release 22.07 (build 41737377)
Triton Server Version 2.24.0
...
I0123 23:54:20.718021 1 pinned_memory_manager.cc:240] Pinned memory pool is created at '0x7f0550000000' with size 268435456
I0123 23:54:20.719918 1 cuda_memory_manager.cc:105] CUDA memory pool is created on device 0 with size 67108864
I0123 23:54:20.728343 1 model_repository_manager.cc:1206] loading: trt_bert:1
I0123 23:54:22.278029 1 libtorch.cc:1917] TRITONBACKEND_Initialize: pytorch
I0123 23:54:22.278058 1 libtorch.cc:1927] Triton TRITONBACKEND API version: 1.10
I0123 23:54:22.279361 1 libtorch.cc:1933] 'pytorch' TRITONBACKEND API version: 1.10
I0123 23:54:22.279422 1 libtorch.cc:1966] TRITONBACKEND_ModelInitialize: trt_bert (version 1)
W0123 23:54:22.283474 1 libtorch.cc:262] skipping model configuration auto-complete for 'trt_bert': not supported for pytorch backend
I0123 23:54:22.284921 1 libtorch.cc:291] Optimized execution is enabled for model instance 'trt_bert'
I0123 23:54:22.284939 1 libtorch.cc:310] Cache Cleaning is disabled for model instance 'trt_bert'
I0123 23:54:22.284946 1 libtorch.cc:327] Inference Mode is disabled for model instance 'trt_bert'
I0123 23:54:22.284955 1 libtorch.cc:422] NvFuser is not specified for model instance 'trt_bert'
I0123 23:54:22.285836 1 libtorch.cc:2010] TRITONBACKEND_ModelInstanceInitialize: trt_bert_0_0 (GPU device 0)
I0123 23:54:27.625852 1 libtorch.cc:2010] TRITONBACKEND_ModelInstanceInitialize: trt_bert_0_1 (GPU device 0)
I0123 23:54:29.381583 1 libtorch.cc:2010] TRITONBACKEND_ModelInstanceInitialize: trt_bert_0_2 (GPU device 0)
I0123 23:54:31.131350 1 libtorch.cc:2010] TRITONBACKEND_ModelInstanceInitialize: trt_bert_0_3 (GPU device 0)
I0123 23:54:32.880875 1 libtorch.cc:2010] TRITONBACKEND_ModelInstanceInitialize: trt_bert_0_4 (GPU device 0)
I0123 23:54:34.630382 1 libtorch.cc:2010] TRITONBACKEND_ModelInstanceInitialize: trt_bert_0_5 (GPU device 0)
I0123 23:54:36.380153 1 libtorch.cc:2010] TRITONBACKEND_ModelInstanceInitialize: trt_bert_0_6 (GPU device 0)
I0123 23:54:38.132072 1 libtorch.cc:2010] TRITONBACKEND_ModelInstanceInitialize: trt_bert_0_7 (GPU device 0)
I0123 23:54:39.884533 1 model_repository_manager.cc:1352] successfully loaded 'trt_bert' version 1
I0123 23:54:39.884634 1 server.cc:559]
+------------------+------+
| Repository Agent | Path |
+------------------+------+
+------------------+------+

I0123 23:54:39.884698 1 server.cc:586]
+---------+---------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------+
| Backend | Path                                                    | Config                                                                                   |
+---------+---------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------+
| pytorch | /opt/tritonserver/backends/pytorch/libtriton_pytorch.so | {"cmdline":{"auto-complete-config":"true","min-compute-capability":"6.000000","backend-directory":"/opt/tritonserver/backends","default-max-batch-size":"4"}} |
+---------+---------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------+

I0123 23:54:39.885569 1 server.cc:629]
+----------+---------+--------+
| Model    | Version | Status |
+----------+---------+--------+
| trt_bert | 1       | READY  |
+----------+---------+--------+

I0123 23:54:39.912066 1 metrics.cc:650] Collecting metrics for GPU 0: Tesla T4
I0123 23:54:39.912400 1 tritonserver.cc:2176]
+----------------------------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| Option                           | Value                                                                                 |
+----------------------------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| server_id                        | triton                                                                                 |
| server_version                   | 2.24.0                                                                                 |
| server_extensions                | classification sequence model_repository model_repository(unload_dependents) schedule_policy model_configuration system_shared_memory cuda_shared_memory binary_tensor_data statistics trace |
| model_repository_path[0]         | /models                                                                                 |
| model_control_mode               | MODE_NONE                                                                                 |
| strict_model_config              | 0                                                                                 |
| rate_limit                       | OFF                                                                                 |
| pinned_memory_pool_byte_size     | 268435456                                                                                 |
| cuda_memory_pool_byte_size{0}    | 67108864                                                                                 |
| response_cache_byte_size         | 0                                                                                 |
| min_supported_compute_capability | 6.0                                                                                 |
| strict_readiness                 | 1                                                                                 |
| exit_timeout                     | 30                                                                                 |
+----------------------------------+----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+

I0123 23:54:39.913476 1 grpc_server.cc:4608] Started GRPCInferenceService at 0.0.0.0:8001
I0123 23:54:39.913767 1 http_server.cc:3312] Started HTTPService at 0.0.0.0:8000
I0123 23:54:39.954918 1 http_server.cc:178] Started Metrics Service at 0.0.0.0:8002
$

BERT での推論

推論スクリプトの invoke_trt_bert.py です。先程の invoke_ft_bert.py との差分は main() 関数だけなので、そこだけ抜粋してます。

...

def main():
  print(f"Loading model from {args.model_path} to BertForTokenClassification.")
  model = BertForTokenClassification.from_pretrained(args.model_path, torchscript=True)
  model.to(device).eval()

  print(f"Loading tokenizer from {args.model_path} to BertForTokenClassification.")
  tokenizer = AutoTokenizer.from_pretrained(args.model_path)

  inputs = ["ジョー・バイデンはアメリカの大統領です。"]

  print(f"input texts : {inputs}")
  print(f"Building input features.")
  input_ids, input_mask, segment_ids, labeled_positions, sub_tokens, tokens = generate_input(inputs,
                                                                                tokenizer, max_len=40)

  batch_size = len(input_ids)
  max_len = 40
  actual_len = len(input_ids[0])

  inputs = []
  outputs = []

  inputs.append(grpcclient.InferInput('input_ids__0', [batch_size, actual_len], "INT32"))
  inputs.append(grpcclient.InferInput('attention_mask__1', [batch_size, actual_len], "INT32"))

  num_labels = 45

  inputs[0].set_data_from_numpy(input_ids)
  inputs[1].set_data_from_numpy(input_mask)

  outputs.append(grpcclient.InferRequestedOutput('probs__0'))

  model_name = "trt_bert"
  results = triton_client.infer(model_name=model_name, inputs=inputs,
                                model_version="1", outputs=outputs)

  probs = results.as_numpy('probs__0')
  probs = probs[:, :max_len * num_labels].reshape(batch_size, max_len, -1) # ★

  probs = torch.tensor(probs)
  labeled_positions = torch.tensor(labeled_positions)

  labeled_positions = labeled_positions.unsqueeze(-1)
  index = torch.zeros(list(labeled_positions.shape[:2]) + [probs.shape[-1]])
  index = (index + labeled_positions).to(torch.int64)
  probs = torch.gather(input=probs, dim=1, index=index)
  preds = torch.argmax(probs, axis=2)
  probs = probs.numpy()
  preds = preds.numpy()

  results = build_result(tokens, sub_tokens, preds, probs, model.config.id2label)
  pp.pprint(results)

...

Triton 側からは flatten してパディングした probs を返すように仕込んだので、 ★のところで末尾を切り飛ばして整形した上で、gather() と argmax() しました。

推論してみます。結果は先ほどと同じですね。

$ docker run --rm --link triton -v `pwd`:/work container-a1 /work/invoke_trt_bert.py
...
Loading model from /work/hf-models/bert-ner to BertForTokenClassification.
Loading tokenizer from /work/hf-models/bert-ner to BertForTokenClassification.
input texts : ['ジョー・バイデンはアメリカの大統領です。']
Building input features.
[ [ {'end': 8, 'start': 0, 'text': 'ジョー・バイデン', 'type': 'PERSON'},
    {'end': 13, 'start': 9, 'text': 'アメリカ', 'type': 'GPE'}]]

長くなってきましたが、あと少しです。次は Neuron で推論してみましょう。

13. Neuron での推論

今度は inf1.xlarge インスタンスで作業を行います。 /home/ssm-user/hf-models が変換元の transformers のモデルを配置したディレクトリです。 inf1.2xlarge で変換したモデルもコピーしてきて、以下の状況です。

$ pwd
/home/ssm-user
$ ls
hf-models  neuron_bert.pt  neuron_t5

それでは推論してみましょう。

BERT の推論

今回は Neuron の推論に Triton を使わないことにしたので、transformers を普通に使うのとあまり差がありません。推論スクリプト invoke_neuron_bert.py は以下のとおりです。省略した関数は Torch-TensorRT のと同じです。

#!/usr/bin/env python
# coding: utf-8

import json
import re
import numpy as np
import torch
import torch.neuron
import datetime
from dataclasses import dataclass, field
from typing import Optional
from transformers import AutoTokenizer, HfArgumentParser
import pprint

pp = pprint.PrettyPrinter(indent=2)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

@dataclass
class Arguments:
    host : Optional[str] = field(default='triton:8001')
    model_path : Optional[str] = field(default='/work/hf-models/bert-ner')
    torch_script_path : Optional[str] = field(default='/work/neuron_bert.pt')

parser = HfArgumentParser((Arguments))
args = parser.parse_args_into_dataclasses()[0]

def joint_sub_word_tokens(tokens):
  #省略

def build_words_to_tokens(sub_tokens):
  #省略

def build_labeld_positions(words_to_tokens):
  #省略

def build_features(text, tokenizer, max_len):
  #省略

def generate_input(texts, tokenizer, max_len=128):
  #省略

def get_span_labels(sentence_tags, inv_label_mapping=None):
  #省略

def get_entities(span_labels, tokens, probs):
  #省略

def build_result(all_tokens, all_sub_tokens, all_predicts, all_probs, id2label):
  #省略

def main():
  texts = ["ジョー・バイデンはアメリカの大統領です。"]
  print(f"input texts : {texts}")

  print(f"Loading tokenizer from {args.model_path} to BertForTokenClassification.")
  tokenizer = AutoTokenizer.from_pretrained(args.model_path)

  print(f"Loading id2label from {args.model_path}.")
  with open(f"{args.model_path}/config.json") as f:
    config = json.load(f)
    id2label = {int(k):v for k, v in config["id2label"].items()}

  print(f"Building input features.")
  input_ids, input_mask, segment_ids, labeled_positions, sub_tokens, tokens = generate_input(texts,
                                                                                tokenizer, max_len=40)

  print(f"Loading Neuron TorchScript from {args.torch_script_path}.")
  model = torch.jit.load(args.torch_script_path)

  batch_size = len(input_ids)
  input_ids = torch.tensor(input_ids).to(device).to(torch.int64)
  attention_mask = torch.tensor(input_mask).to(device).to(torch.int32)
  labeled_positions = torch.tensor(labeled_positions).to(device).to(torch.int32)
  num_labels = torch.tensor([45] * batch_size).to(device).to(torch.int32)

  preds, probs = model(input_ids, attention_mask, labeled_positions, num_labels)

  probs = probs.cpu().detach().numpy()
  preds = preds.cpu().detach().numpy()

  results = build_result(tokens, sub_tokens, preds, probs, id2label)
  pp.pprint(results)

if __name__ == "__main__":
  with torch.no_grad():
    main()

Neuron で推論する場合は –e と –device オプションを以下のように指定します。

–device はコンテナから Neuron ドライバにアクセスできるようにする為の設定です。
–e の NEURON_RT_NUM_CORES はプロセスに Neuron コアを何個割り当てるかの指定です³³。

$ docker run --rm -e "NEURON_RT_NUM_CORES=1" --device "/dev/neuron0:/dev/neuron0" \
    -v `pwd`:/work container-c /work/invoke_neuron_bert.py
/opt/conda/lib/python3.7/site-packages/transformers/tokenization_utils_base.py:2307: FutureWarning: The `pad_to_max_length` argument is deprecated and will be removed in a future version, use `padding=True` or `padding='longest'` to pad to the longest sequence in the batch, or use `padding='max_length'` to pad to a max length. In this case, you can give a specific length with `max_length` (e.g. `max_length=45`) or leave max_lengthto None to pad to the maximal input size of the model (e.g. 512 for Bert).  FutureWarning,
input texts : ['ジョー・バイデンはアメリカの大統領です。']Loading tokenizer from /work/hf-models/bert-ner to BertForTokenClassification.Loading id2label from /work/hf-models/bert-ner.Building input features.
Loading Neuron TorchScript from /work/neuron_bert.pt.
[ [ {'end': 8, 'start': 0, 'text': 'ジョー・バイデン', 'type': 'PERSON'},
    {'end': 13, 'start': 9, 'text': 'アメリカ', 'type': 'GPE'}]]

相変わらず CTRL+C で止めてるのですが無駄に長くなるので記載してません。最後は T5 です。

T5 の推論

推論スクリプト invoke_neuron_t5.py は以下のとおりです。省略したところは変換に使用したスクリプトと同じ内容になります。

#!/usr/bin/env python
# coding: utf-8

import os
import numpy as np
import torch
from torch.nn import functional as F
import datetime
from dataclasses import dataclass, field
from typing import Optional, Any
from transformers import PreTrainedModel, T5Tokenizer, HfArgumentParser, T5Config, T5ForConditionalGeneration
from transformers.generation_utils import GenerationMixin
from transformers.modeling_outputs import BaseModelOutput, Seq2SeqLMOutput
import torch.neuron

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

@dataclass
class Arguments:
    host : Optional[str] = field(default='triton:8001')
    model_path : Optional[str] = field(default='/work/hf-models/t5-1.0-base')
    torch_script_path : Optional[str] = field(default='/work/neuron_t5')

parser = HfArgumentParser((Arguments))
args = parser.parse_args_into_dataclasses()[0]

def reduce(hidden: torch.Tensor, index: int):
  #省略

class NeuronEncoder(torch.nn.Module):
  #省略

class NeuronDecoder(torch.nn.Module):
  #省略

class NeuronGeneration(PreTrainedModel, GenerationMixin):
  #省略

def tokenize(tokenizer, text, max_len):
  return tokenizer(text, return_tensors='pt', truncation=True, max_length=max_len,
            padding="max_length")

def decode(tokenizer, text, result):
  output = result.as_numpy("output_ids")
  ft_output_len = result.as_numpy("sequence_length")
  return tokenizer.decode(output[0][0][:ft_output_len[0][0]], skip_special_tokens=True)

 def main(tokenize, text):

  print(f"input text: {text}")

  print(f"Loading tokenizer from {args.model_path}")
  tokenizer = T5Tokenizer.from_pretrained(args.model_path)
  bos_token_id = tokenizer.bos_token_id if tokenizer.bos_token_id else 0
  eos_token_id = tokenizer.eos_token_id if tokenizer.eos_token_id else 1
  print(f"bos_token_id = {bos_token_id}")
  print(f"eos_token_id = {eos_token_id}")

  print(f"Loading model from {args.torch_script_path}")
  model = NeuronGeneration.from_pretrained(args.torch_script_path)

  print(f"Tokenizing input.")
  input_tokens = tokenize(tokenizer, text, max_len=40)

  input_ids = input_tokens.input_ids
  attention_mask = input_tokens.attention_mask

  gen_kwargs = {}
  gen_kwargs["attention_mask"] = attention_mask
  gen_kwargs["bos_token_id"] = bos_token_id
  gen_kwargs["eos_token_id"] = eos_token_id
  gen_kwargs["top_k"] = 1
  gen_kwargs["top_p"] = 1.0
  gen_kwargs["temperature"] = 1.0
  gen_kwargs["max_length"] = 40
  gen_kwargs["num_beams"] = 1
  gen_kwargs["length_penalty"] = 1.0
  gen_kwargs["repetition_penalty"] = 1.0

  with torch.no_grad():
    result = model.generate(input_ids, **gen_kwargs)

  print(f"Decoding output.")
  print(tokenizer.decode(result[0], skip_special_tokens=True))

if __name__ == "__main__":
  with torch.no_grad():
    main(tokenize, "その男は窃盗犯に問われた。")

以下のようにして実行します。

$ docker run --rm -e "NEURON_RT_NUM_CORES=1" --device "/dev/neuron0:/dev/neuron0" \
    -v `pwd`:/work container-c /work/invoke_neuron_t5.py
input text: その男は窃盗犯に問われた。
Loading tokenizer from /work/hf-models/t5-1.0-base
bos_token_id = 0
eos_token_id = 1
Loading model from /work/neuron_t5
Tokenizing input.
Decoding output.
その男は盗んだ人に取られた。

若干変換が微妙ですが、ちゃんと推論出来ているようです。

ここからは、変換した各種モデルに負荷をかけてスループットを計測してみたので、その結果を紹介したいと思います。

14. 計測結果

ここまでで変換してきたモデルを Web API に仕立てて負荷試験しました。負荷生成は以下のように行っています。

負荷生成は t3.medium の別ノードから locust を用いて行った。
計測前に 20 秒のウォームアップを実施。
負荷計測時間は 480 秒
ユーザは 30 秒毎に 16 ユーザまでランプアップ
各ユーザは待ち時間なしで Web API をコールし続ける。
Web API への入力テキストはデータセットから 300 件抽出して呼び出し毎にランダム選択した。

また以降の表記で共通の説明として、

nopad : 入力へのパディングなし。
padx8 : 入力に対しシーケンス長が8の倍数になるようにパディング。

となります。まずは BERT の計測結果です。

BERT の計測結果

スループットを比較すると以下のようになりました。

FasterT
- transformers のモデルを FasterTransformer に変換して Triton にデプロイし、gunicorn/Falcon のフロントエンドを被せた形式。
- gunicorn の設定 : workers = 8
- Triton の設定 : max_batchsize=16, dynamic_batching {}
- 0.4 % 程度の確率で内部エラー 500 が発生。数分で Triton が落ちる現象が発生したので計測時間を 30 秒とした値。なので参考値扱いです。
Transf
- gunicorn/Falcon の worker が transformers のモデルを属性として保持する形式。
- gunicorn には次を設定 :workers = 4
Neuron
- gunicorn/Falcon の worker が transformers のモデルを Neuron にコンパイルした TorchScript のモデルを属性として保持する形式。
- gunicorn の設定 : workers = 4
- 環境変数 NEURON_RT_NUM_CORES=1 で各ワーカに Neuron Core を一つずつ割り当て
- バッチサイズ = 1, 入力はパディングして固定シーケンス長 = 40 とした。
TensorRT
- transformers のモデルを TorchScript にコンバートした上で、Torch-TensorRT でコンパイルしたものを Triton にデプロイし、gunicorn/Falcon のフロントエンドを被せた形式。
- gunicorn の設定 : workers = 8
- Triton の設定 : max_batchsize=0, instance_group [ { count: 8 } ]
- バッチサイズ = 1, 入力はパディングして固定シーケンス長 = 40 とした。
TorchS
- transformers のモデルを TorchScript に変換して Triton にデプロイし、gunicorn/Falcon のフロントエンドを被せた形式。
- gunicorn の設定 : workers = 8
- Triton の設定 : max_batchsize=8, instance_group [ { count: 8 } ], dynamic_batching {}

結果として Torch-TensorRT が 341.50 req/sec で最良になりました。

ただ、期待したほど速くなっていませんね。。。リソースの使用状況を確認すると以下のとおりでかなり余裕があるように見えます。

ですが、Triton の Model Analyzer³⁴ を用いて Triton 単体で計測すると最大 461 req/sec で、この状態でも GPU を使いきれてませんでした。バッチサイズ = 1 で固定かつシーケンス長 = 40 の比較的短いテキストではリソースを使い切るのは難しいのかも知れませんね。もっと長いテキストで大きなバッチを組めれば、より効果が望める気がします。今後 Triton の Torch-TensorRT で動的バッチングが使えるようになったらリトライしたいですね。

また、Neuron は 255.10 req/sec で Torch-TensorRT に及びませんでしたが、インスタンスの単価を考えると推論辺りのコストでは最良になりました。

それでは次に T5 を見ていきましょう。

T5 の計測結果

スループットを比較すると以下のようになりました。

FasterT
- transformers のモデルを FasterTransformer にコンバートして Triton にデプロイし、gunicorn/Falcon のフロントエンドを被せた形式。
- gunicorn の設定 : workers = 8
- Triton の設定 : max_batchsize=16, dynamic_batching {}
- 入力、出力最大シーケンス長は共に 40
Transf
- gunicorn/Falcon の worker が transformers のモデルを属性として保持する形式。
- gunicorn の設定 : workers = 4
- 入力、出力最大シーケンス長は共に 40
Neuron
- gunicorn/Falcon の worker が Neuron のモデルを属性として保持する形式。
- 入力、出力シーケンス長は共に 40 固定。
全てのモデルでテキスト生成は greedy で実施。

FasterTransformer で入力を 8 の倍数にパディングしたケースが 76.80 req/sec で最良でした。パディングなしとの差は Triton の動的バッチングによるものかもしれません。

transformers は fp16 での推論が fp32 よりも遅くなりました。この原因はよくわかりません。そう言えば第14回でも T5 の fp16 でガッカリしてたような。。。

Neuron は transformers と大差ない結果がでました。これもインスタンスの単価を考えればアリかという気がしなくもないですが、シーケンス長が固定なので出力長が長くなると苦しくなるでしょう。 T5 1.1 を使えないのも減点ですね。

最後に GPT です。

GPT の計測結果

スループットを比較すると以下のようになりました。

FasterT
- transformers のモデルを FasterTransformer にコンバートして triton にデプロイし、gunicorn/Falcon のフロントエンドを被せた形式。
- gunicorn の設定 : workers = 8
- Triton の設定 : max_batchsize=16, dynamic_batching {}
- 入力最大シーケンス長は 40、出力シーケンス長は 40 固定。
- パディング時は入力の右側に付加。
Transf
- gunicorn/Falcon の worker が transformers のモデルを属性として保持する形式。
- gunicorn の設定 : workers = 4
- 最大シーケンス長は 80（そのうち入力分で最大40）
- パディング時は入力の左側に付加。generate() には attention_mask を付けて投入。
プロンプトとして“幸運にも彼は死ななかった。::” を投入すると、"幸運にも彼は死ななかった。::運の良いことにも彼は死ななかった。"を生成するようファインチューニングした。
全てのモデルでテキスト生成は greedy で実施。

こちらも T5 と似た結果で FasterTransformer の 8 の倍数パディングが 39.40 req/sec で最良でした。これだけ差があるとテキスト生成モデルの推論には FasterTransformer を使いたくなりますね。

FasterTransformer で「GPT への入力の右側にパディングする」というのが意味不明な感じがしますが、有効なトークン数もセットで FasterTransformer に渡すので、そこはうまくやってくれるのでしょう。パディングの意味合いとしては、長さを揃えることで動的バッチングが効きやすくなるのかと思います。

あと、gunicorn で workers = 8 ですからバッチサイズが > 8 になる訳がないのに Triton 側で max_batchsize=16 にしてるので、Triton 側で無駄な待ちを発生させてしまっていますね。他にもセッティングとして詰め切れてないところは多々あると思うので、そのつもりで見てもらえれば。

逆に transformers の場合は左側にパディングをして attention_mask も添えて generate() する必要がありました。こちらは Triton がないので動的バッチングが効きません。パディング有無であまり差がないのも納得ですね。

最後に FasterTransformer 限定でテキスト生成モデル間での比較です。

テキスト生成モデル間での比較

ここまで説明してきた T5, GPT に加え、りんなさんが公開している GPT-1B モデル³⁵も加えて比較してみます。

GPT-1B も GPT と同様のファインチューニングして FasterTransformer に変換、8 の倍数パディングでの結果です。さすがに重たいですねぇ。。。これでも素の transformers と比べると 3 倍速かったんですよ。

15. おわりに

今回は、transformers のモデルを様々な手法で変換して効果を確認してみました。今回は手を出しませんでしたが機会があれば int8 への量子化とかも試してみたいです。次回は久々に Rasa³⁶ の話にしようかと思います。transformers と組み合わせてみようかなと。

https://github.com/huggingface/transformers ↩
https://github.com/NVIDIA/FasterTransformer ↩
https://github.com/pytorch/TensorRT ↩
https://github.com/aws-neuron/aws-neuron-sdk ↩
https://aws.amazon.com/jp/machine-learning/inferentia/ ↩
Web API 化の巧拙や個々のパターンをあまり突き詰められてない所もありますが、この連載の性能比較はプロレス的な感覚で楽しんで頂けると幸いです。 ↩
https://gunicorn.org/ ↩
別に Falcon である理由はないです。過去の作業を流用したかったというくらいですね。 https://falcon.readthedocs.io/en/stable/ ↩
Triton は AWS Inferentia に対応しているので、パターン 2 の構成にできたと思うのですが、今回そこまで手が回りませんでした。本当は FasterTransformer だけ試して終わりにするつもりだったんですが、どんどん広がっていって。。。 ↩
https://github.com/triton-inference-server/server ↩
https://github.com/tensorflow/serving ↩
https://github.com/triton-inference-server/server/blob/main/docs/user_guide/model_repository.md ↩
https://github.com/triton-inference-server/server/blob/main/docs/user_guide/model_repository.md#model-files ↩
https://github.com/NVIDIA/FasterTransformer#support-matrix ↩
https://developer.nvidia.com/blog/optimizing-t5-and-gpt-2-for-real-time-inference-with-tensorrt/ ↩
https://github.com/triton-inference-server/fastertransformer_backend ↩
自己回帰処理まで含んだ形の SavedModel なのでデプロイはできるのですが、入出力が固定長になったりビーム探索や Top K サンプリングのパラメータを変える為に SavedModel を作り直さないといけなかったり。 ↩
https://developer.nvidia.com/tensorrt ↩
https://github.com/pytorch/TensorRT/releases/tag/v1.3.0 ↩
https://docs.nvidia.com/deeplearning/frameworks/support-matrix/index.html ↩
https://aws.amazon.com/jp/machine-learning/inferentia/ ↩
https://aws.amazon.com/jp/machine-learning/trainium/ ↩
https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/arch/model-architecture-fit.html#aws-inferentia-neuroncore-v1 ↩
https://github.com/aws-neuron/aws-neuron-sdk/issues/443 ↩
https://github.com/aws-neuron/aws-neuron-sdk/issues/427 ↩
https://nvidia.custhelp.com/app/answers/detail/a_id/5323/~/known-issue%3A-nvidia-vgpu-software-graphics-driver-installation-fails-on-some ↩
Neuron のコンパイルに Inferentia は不要だったかもしれませんが、今回は話を簡単にするために Inf1 インスタンスで行いました。inf1.2xlarge なのは inf1.xlarge だと物理メモリが足りなくて OOM Killer が発動したから（だったような。。。） ↩
ECR に Docker イメージを push するのに必要だと思うので、今回は不要だったかもしれません。。。 ↩
22.07 は NVIDIA が配布するコンテナイメージのバージョンで Triton 自体は 2.24.0 になります²⁰。 ↩
https://github.com/NVIDIA/FasterTransformer/commit/c2547131bb058683aa0ab34e1f0ad3a6d34c67c5 ↩
処理的にはパディングを付けてますが、変換自体は max_num_labels = 45 として学習データの分類数(45)と一致させる形で記述してます。 ↩
白状すると「TorchScript に変換して Triton にデプロイする時に返却するテンソルのシェイプを動的にしたかった。」という理由で書いたコードが Neuron 版にも残っちゃったという話で（すみません）。。。 ↩
inf1.xlarge には Neuron コアが 4 個あるはずなのですが、4 を指定したらエラーになりました。。。 ↩
https://github.com/triton-inference-server/model_analyzer ↩
https://huggingface.co/rinna/japanese-gpt-1b ↩
https://github.com/RasaHQ/rasa ↩