Amazon SageMaker で Hugging Face Transformers を使用して GPT-J 6B をデプロイする

EleutherAI の GPT-J 6B（オープンソースの 60 億パラメータ言語モデル）を、Amazon SageMaker と Hugging Face Inference Toolkit を使用してリアルタイム推論用にデプロイする方法が Hugging Face によって詳細に説明されています。このアプローチは、Amazon SageMaker によって課される 60 秒のリクエスト制限を超える可能性がある、モデルのロード時間の長さという重大な課題を解決します。

本番環境向けのモデルロード時間の最適化

GPT-J 6B を本番環境にデプロイすることは、そのメモリ使用量とロード速度のため、困難です。60 億パラメータのモデルの重みは、約 24GB のメモリを表します。モデルを float32 でロードするには、少なくとも 48GB の CPU RAM（初期重み用とチェックポイント用でそれぞれ 1 つ）が必要です。

アクセシビリティを向上させるために、EleutherAI は float16 の重みを提供しています。メモリ使用量を削減するための transformers オプションと組み合わせると、CPU RAM の要件は約 12.1GB に減少します。しかし、標準的なロード方法は依然として低速です。

標準的なロード: P3.2xlarge AWS EC2 インスタンスでモデルをロードするには、約 3 分 32 秒かかります。
ディスク保存されたロード: モデルをディスクに保存すると、これを 1 分 23 秒に短縮できます。

Amazon SageMaker はリクエストへの応答に 60 秒の制限があるため、これらのロード時間は、モデルをスケーラブルで信頼性の高い本番環境のワークロードには不向きであることを意味します。

torch.save によるロード時間の加速

推奨される from_pretrained メソッドの代わりに torch.save(model, PATH) と torch.load(PATH) を使用することで、GPT-J のロード時間を 1 分 23 秒から 7.7 秒に短縮できます。これは、約 10.5 倍の高速化です。

重要要件: 互換性を避けるために、ユーザーはモデルを保存するときに使用した PyTorch と Transformers のバージョンを、ロードするときに使用するものと一致させる必要があります。

Amazon SageMaker でのデプロイメントワークフロー

GPT-J 6B をリアルタイム推論用にデプロイするには、以下のワークフローを使用します。

モデルのシリアライズ: from_pretrained を使用して GPT-J をロードし、torch.save() を使用して .pt ファイルとして保存します。
アーティファクトの作成: モデルの重みと必要なファイル（tokenizer.json など）を含む model.tar.gz アーカイブを作成します。このアーティファクトは、S3 バケットにアップロードされます。
エンドポイントのデプロイ: Amazon SageMaker Python SDK の HuggingFaceModel クラスを使用してモデルをデプロイします。

このデプロイメントでは、Hugging Face は ml.g4dn.xlarge インスタンスタイプ（NVIDIA T4 GPU）を推奨しており、コストは約 $500/月です。

推論と使用のベストプラクティス

エンドポイントがデプロイされたら、予測は predictor.predict メソッドを介して実行されます。Hugging Face Inference toolkit は、リクエストペイロード内の parameters 属性を通じて、生成プロセスをカスタマイズできます。

生成戦略

Greedy Search: デフォルトのリクエスト方法です。最初のリクエストの後、推論時間は約 3 秒です。
Beam Search: パラメータに `