optimum-neuron を使用した AWS Inferentia2 上への Llama 2 のデプロイ

Hugging Face は optimum-neuron を AWS Neuron SDK と統合し、AWS Inferentia2 アクセラレータ上での Llama 2 モデルのデプロイを可能にしました。この統合により、ユーザーは専用ハードウェアを活用して、エンコーディング時間、レイテンシ、スループットの面でテキスト生成パフォーマンスを向上させることができます。

デプロイのワークフローとセットアップ

AWS Inferentia2 に Llama 2 をデプロイするには、特定の環境セットアップと、モデルを Neuron デバイスと互換性のあるシリアル化された形式に変換するためのモデルコンパイルステップが必要です。

環境構成

ユーザーは、主に以下の3つの方法で Inferentia2 インスタンスをセットアップできます：

Hugging Face Neuron Deep Learning AMI (DLAMI): 推奨される方法であり、Optimum Neuron、Neuron Drivers、Transformers、Datasets、Accelerate を含むパッケージ化されたライブラリが提供されます。
Hugging Face Neuron SDK DLC: Amazon SageMaker でのデプロイに使用されます。
手動インストール: 新しいインスタンスで optimum-neuron のインストール手順に従います。

モデルのエクスポートとコンパイル

Neuron デバイスは静的な形状（static shapes）を必要とするため、実行前にモデルをコンパイルする必要があります。NeuronModelForCausalLM API を使用すると、エクスポート時に以下のパラメータを指定できます：

Compiler Arguments: コア数（各 Neuron デバイスには2つのコアがあります）と精度（例：float16）を定義します。
Input Shapes: batch_size と sequence_length の静的な次元を設定します。sequence_length は、入力コンテキスト、KV キャッシュ、および最大出力長を制限するため、非常に重要です。

コンパイルが完了すると、モデルをローカルに保存したり、再利用のために Hugging Face Hub にプッシュしたりできます。

テキスト生成機能

optimum-neuron は、transformers ライブラリを使用した標準的なテキスト生成、または簡略化された optimum-neuron パイプラインをサポートしています。

生成戦略

サポートされている生成戦略には以下が含まれます：

Greedy search
top-k および top-p を使用した multinomial sampling（temperature を含む）
repetition penalty などのほとんどの logits 前処理フィルタ

実装オプション

効率的なデプロイのために、optimum-neuron の pipeline API を使用すると、Hub からコンパイル済みモデルをロードし、単一の関数呼び出しでテキストを生成できます。

パフォーマンス・ベンチマーク

ベンチマークは、inf2.xlarge（予算重視モデル用）および inf2.48xlarge（レイテンシおよびスループット最適化モデル用）において、Llama 2 7B および 13B モデルを使用して様々な構成で行われました。すべてのモデルは最大シーケンス長 2048 を使用しました。

エンコーディング時間

エンコーディング時間（入力トークンを処理して最初の出力トークンを生成するまでの時間）は、ユーザーが感じるレイテンシの重要な指標です。256 入力トークン（一般的な Q&A）の場合、エンコーディング時間は 0.3s (Llama2 7B-B) から 0.9s (Llama2 7B-T) の範囲でした。768 入力トークン（一般的な RAG）の場合、時間は 0.5s (Llama2 7B-B) から 5.2s (Llama2 13B-T) の範囲でした。

エンドツーエンドのレイテンシ

エンドツーエンドのレイテンシは、シーケンス長が 1024 トークンに達するまでの総時間を測定します。ハイエンドの inf2.48xlarge インスタンスでは、Llama2 7B-L は 768 新規トークンに対して 6.2s のレイテンシを達成し、Llama2 13B-L は 10.2s を要しました。「予算重視」の inf2.xlarge 上のモデルは、768 新規トークンに対して 47.3s に達し、著しく高いレイテンシを示しました。

スループット

スループットは、1秒あたりの総トークン数（エンドツーエンドのレイテンシを batch_size * sequence_length で割ったもの）として計算されます。

ハイパフォーマンス: Llama2 7B-T 構成は、256 新規トークンに対して最大 750 tokens/second を達成しました。
予算重視: Llama2 7B-B モデルは 22 から 32 tokens/second を達成しました。Hugging Face は、平均的な人間の読解速度を考慮すると、これはストリーミングのユースケースには十分であると述べています。

技術的な制限と今後の課題

パフォーマンスは強力ですが、Hugging Face は改善が必要な2つの主要な領域を特定しています：

スループットのスケーリング: 現在、スループットは主にバッチサイズを増やすことで向上しますが、これはデバイスメモリによって制限されます。代替案としてパイプライニングの統合が進められています。
コンテキスト長: 静的なシーケンス長の要件により、非常に長いコンテキストを扱う能力が制限されています。チームは、潜在的な解決策として attention sinks を検討しています。