こんにちは。日経電子版開発でバックエンドを担当している田中です。
本投稿では、Google Cloud の Cloud Run と、NVIDIA で開発されている Triton Inference Server と Huggingface を使って、文を受け取って埋め込みを返す API を作る方法を検証したので紹介させていただきます。
なお本投稿では細かい用語などの説明やコンソールなどでの詳細な操作は省かせていただきます。
動機
サーバーレスで、それなりのマシンリソースを使えて、かつ手っ取り早く API を構築する方法を模索していたためです。現状サービスや社内ツールなどで使う予定はないのですが、マシンリソースが要求されるようなモデルをとりあえずで動かす手段を用意しておきたくて調査しました。特に社内ツール作成に推論 API が必要になった時、大半の時間は稼働しないが計算するときはそれなりに重たいリソースが必要…となりやすいので、どうしても費用対効果の面での懸念がありました。今回はコストや機能に柔軟性をもちつつも、一度作ればテンプレ的に取り回しが効きそうに見えたため、表題の組み合わせを検証することにしました。
他サービスとの比較でいえば、AWS Sagemaker Serverless Inference がやりたいことを叶えてくれそうなサービスですが、メモリが 6GB しか使えず、昨今のモデルを動かすのには不十分でした。その点 Cloud Run であればメモリは 32GB、vCPU が 8 まで利用できるので、BERT ぐらいの規模なら実行可能です。
Triton Inference Serverについて
Cloud Run や Huggingface に関しては言わずと知れたツールかと思いますので、 Triton Inference Server についてのみ簡単に説明させていただきます。
Triton Inference Server（以下 Triton）は NVIDIA で開発されている推論サーバー構築に利用できるコンテナです。基本機能として、モデルバイナリと、エンドポイントのスキーマや設定を記述したファイルを所定の場所に配置するだけで、 REST や gRPC に対応した推論 API が完成します。高いスループットを謳っており、こちらのブログでも高パフォーマンスであることが検証されています。
計算には GPU、CPU どちらも利用できます。また、AWS で提供されている推論に最適化された CPU を搭載する Inferentia インスタンスでも実行可能です。
モデルのライブラリも Pytorch や Tensorflow、ランタイムも ONNX など各ランタイムに対応しており、Python スクリプトによるアドホックな処理も記述できます。
大まかな手順
今回作成する文埋め込みを返す処理ですが、Huggingface の Tokenizer は ONNX などのランタイム形式に変換することができません。CLIP を使うときはできるようなので、本当はできる気はするのですが、良さそうな方法が出てこなかったので今回はランタイムに変換できないとして進めます。 そのため、文字列を受けとってトークナイズして文埋め込みを推論し、その結果を返す処理を書いた Python スクリプトと一緒にデプロイすることにします。
大まかには次の手順で進めます。
- Triton 用の推論スクリプトとモデルバイナリを作成
- Triton 用の設定を記述
- Google Cloud Storage (GCS) にスクリプトとモデルバイナリを配置
- Triton イメージを作成して Artifact Registry に配置
- Cloud Run で Triton イメージをデプロイ
- リクエストして動作確認
Triton 用の推論スクリプトを作成
モデルは最近話題の multilingual-e5-large を使用します。埋め込み処理は multilingual-e5-large のサンプルで記述されている、各トークンに対応する出力を平均プーリングする方法で計算しています。
また、Triton はテキストをバイトとして受け取るのでトークナイズ前にデコードが必要になります。
これを model.py として保存します。
import triton_python_backend_utils as pb_utils
from torch import nn
from transformers import AutoTokenizer
class HFEmbeddingInference(nn.Module):
def __init__(self):
super().__init__()
self.tokenizer = AutoTokenizer.from_pretrained("intfloat/multilingual-e5-large")
current_dir = os.path.dirname(os.path.realpath(__file__))
self.model = torch.jit.load(os.path.join(current_dir, "multilingual-e5-large.pt"))
self.model.eval()
self.max_length = self.tokenizer.model_max_length
def forward(self, input_text: list[str]):
batch_dict = self.tokenizer(
input_text,
max_length=self.max_length,
padding=True,
truncation=True,
return_tensors="pt",
)
outputs = self.model(
batch_dict["input_ids"],
attention_mask=batch_dict["attention_mask"],
)
logger = pb_utils.Logger
last_hidden_states = outputs[0]
attention_mask = batch_dict["attention_mask"]
last_hidden = last_hidden_states.masked_fill(
~attention_mask[..., None].bool(), 0.0
)
out = last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]
logger.log_info(str(out.shape))
return out
class TritonPythonModel:
def initialize(self, args):
self.model_config = model_config = json.loads(args["model_config"])
# Get OUTPUT0 configuration
output0_config = pb_utils.get_output_config_by_name(model_config, "OUTPUT0")
# Convert Triton types to numpy types
self.output0_dtype = pb_utils.triton_string_to_numpy(
output0_config["data_type"]
)
# Instantiate the PyTorch model
self.embedding_model = HFEmbeddingInference()
def execute(self, requests):
output0_dtype = self.output0_dtype
responses = []
for request in requests:
# Get INPUT0
in_0 = pb_utils.get_input_tensor_by_name(request, "INPUT0")
texts = [text.decode("utf-8") for text in in_0.as_numpy().tolist()]
with torch.inference_mode():
out_0 = self.embedding_model(texts)
out_tensor_0 = pb_utils.Tensor(
"OUTPUT0", out_0.numpy().astype(output0_dtype)
)
inference_response = pb_utils.InferenceResponse(
output_tensors=[out_tensor_0]
)
responses.append(inference_response)
# You should return a list of pb_utils.InferenceResponse. Length
# of this list must match the length of `requests` list.
return responses
def finalize(self):
print("Cleaning up...")
モデル本体は 2GB ほどあり、コンテナ起動時に毎回 Huggingface からのダウンロードが発生するのは重たいので、TorchScript 化して読み込むことにしました。このスクリプト実行時点では multilingual-e5-large.pt の名前で TorchScript 化済みです。このあと GCS にアップロードしますが、モデルを読み込むときこのパスのままでも問題ありませんでした。
Triton 用の設定
複数の文を同時に受け取って、各々の埋め込みを返せるように設定します。設定値には引数名と型と次元を書きます。配列次元が 1 次元のとき、-1 をつけると 1 次元の自由長になるようです。
他にも使う計算リソースの設定を記述でき、今回は CPU なので CPU にします。CPU を使うモデルと GPU を使うモデルを同時にデプロイしつつ、モデルごとに切り替えられるようです。
これを config.pbtxt として保存します。
name: "multilingual_e5_large"
backend: "python"
input [
{
name: "INPUT0"
data_type: TYPE_STRING
dims: [ -1 ]
}
]
output [
{
name: "OUTPUT0"
data_type: TYPE_FP32
dims: [ 1024 ]
}
]
instance_group [{ kind: KIND_CPU }]
GCSにスクリプトとモデルバイナリを配置
バケットを作り、次のような階層で配置をします。1 はバージョンに相当する整数です。
model-repository
`-- multilingual_e5_large
|-- 1
| |-- model.py
| `-- multilingual-e5-large.pt
`-- config.pbtxt
Tritonイメージを作成してArtifact Registryに配置
Triton イメージ自体は配布されていますが、Python スクリプトを利用する場合、Python から transformers と torch を呼び出せるようにインストールされたイメージを別個用意する必要があります。
FROM nvcr.io/nvidia/tritonserver:23.08-py3
ENV PYTHONUNBUFFERED 1
ENV PIP_NO_CACHE_DIR 1
RUN pip3 install \
"transformers==4.33.2" "torch==2.0.1"
なお、公式のイメージはフレームワークやランタイム全部載せになっています。Tensorflow や ONNX など今回使わないランタイムも入っています。今回は torch も別途 pip 経由で入れてしまったため、デフォルトで用意されている LibTorch は必要ありません。
これらをなくして Docker イメージをダイエットしたい時、GitHub レポジトリにビルド用のスクリプトのソースが提供されており、そこから必要なバックエンドだけを入れたイメージと Dockerfile を作ってくれるようです。
また、公式では追加のライブラリを入れる方法は、使いたい Python バージョンがイメージで用意されているものと違うときに行う方法という形で紹介されています。通常通りデフォルトのバージョンで問題ない時はどのようにライブラリを追加するか案内がないので、これがベストプラクティスではない可能性があります。
完成したら Artifact Registry に push します。ローカルでも実行できるので、コンテナを起動すれば動作確認できます。
Cloud RunでTritonイメージをデプロイ
コンソールぽちぽちで OK です。
実行コマンドに次のコマンドを渡します。またポートフォワード先は 8000 にします。スペックはメモリを 32GiB、vCPU を 8 にしました。
tritonserver --model-repository=gs://<YOUR-BUCKET-NAME>/model-repository
リクエストして動作確認
三つの文章を送って、結果を受け取り、cos 類似度行列を出力するスクリプトを組みました。
query: という接頭辞がありますが、これは使っているモデルの仕様です。
import numpy as np
import tritonclient.http as httpclient
from tritonclient.utils import *
model_name = "multilingual_e5_large"
with httpclient.InferenceServerClient("xxxxxxxx.a.run.app", ssl=True) as client:
input0_data = np.array(
[
"query: dummy",
"query: 食べられる薔薇を育てている農家は日本に数件しかありません。",
"query: 日本の薔薇農家",
],
dtype=object,
)
inputs = [
httpclient.InferInput("INPUT0", input0_data.shape, "BYTES"),
]
inputs[0].set_data_from_numpy(input0_data, binary_data=False)
outputs = [
httpclient.InferRequestedOutput("OUTPUT0"),
]
print("requests...")
response = client.infer(model_name, inputs, outputs=outputs)
result = response.get_response()
output0_data = response.as_numpy("OUTPUT0")
output0_data = output0_data / np.linalg.norm(output0_data, axis=1, keepdims=True)
print(output0_data.dot(output0_data.T))
結果
[[1.0000001 0.7383819 0.7169244 ]
[0.7383819 1.0000001 0.91994643]
[0.7169244 0.91994643 1.0000002 ]]
cos 類似度としてみたとき、全部高いのが気になりますが、順序関係としては正しいので、うまく動いていそうです。
なお、コールドスタートがあり、しばらく放置してからリクエストすると受け付けられるようになるのに 2 分程度かかってしまいます。
感想
今回は文を受け取って埋め込みを返す API を作成しました。Cloud Run で設定に詰まるかなと思いましたが、そこは何事もなく進みました。
また、パフォーマンスを追求するなら、文を受け付けられるようにするより、トークナイズした結果を受け付ける方が良さそうです。Python がなくなりランタイムのみで実行できるためです。ただ、 API として見た時に、Tokenizer が使えない言語では呼び出しできなくなってしまうので、汎用性を目指すなら文を受け付けられる形式の方が良さそうです。
またパフォーマンスを調べておらず、 WebAPI フレームワークを使って愚直に作った場合とどちらが速いのかも気になっており、今後時間がある時に調査しておきたいと思います。
また Triton は定期実行バッチも作れるようなのでその実装やパフォーマンスも今後調査していきたいです。
終わりに
テクノロジーメディアを目指す日本経済新聞社ではデジタルサービスにおける ML エンジニア、データサイエンティスト、バックエンドエンジニアを募集しています。自然言語処理の最新技術の検証、アプリケーション開発なども積極的に行なっております。ご興味を持たれた方は是非ご応募ください。
ML エンジニア募集ページはこちら
https://herp.careers/v1/nikkei/pbcSlV_6jPUc
データサイエンティスト・データアナリストの募集ページはこちら
