Hugging Face Text Generation Inference が AWS Inferentia2 をサポート開始

TL;DR

Hugging Face は、Amazon SageMaker を通じて AWS Inferentia2 上で Text Generation Inference (TGI) を一般提供し、GPU 不要で本番レベルの大規模言語モデル（LLM）を低レイテンシかつ高同時実行で提供できるパスを提供します。

Text Generation Inference (TGI) とは？

TGI は LLM 用に特化した推論サーバーで、以下を提供します：

  • 高スループットを実現するテンソル並列化と連続バッチ処理。
  • Llama、Mistral などの人気オープンソースモデルに最適化されたサポート。
  • Grammarly、Uber、Deutsche Telekom などの企業が利用する本番環境向けデプロイ。

Inferentia2 が重要な理由

AWS Inferentia2 チップは GPU よりも低コストで高性能な推論を提供します。TGI と Inferentia2、SageMaker を統合することで、Hugging Face は以下を提供します：

  • シームレスで管理されたデプロイ体験。
  • HuggingChat、OpenAssistant、Hugging Face のサーバーレスエンドポイントを支える同じバックエンド技術との互換性。
  • AWS 固有ハードウェアを好む顧客向けの代替コンピュートスタック。

Inferentia2 上で Zephyr 7B をデプロイする – ステップバイステップガイド

このブログ記事では、ml.inf2.8xlarge インスタンス上に Zephyr 7B モデル（Mistral‑7B‑v0.1 の DPO ファインチューニング版）をデプロイする手順を解説しています。手順は自己完結型で、SageMaker ユーザーであれば誰でも再現可能です。

1. 開発環境のセットアップ

pip install transformers "sagemaker>=2.206.0" --upgrade --quiet

このスクリプトは IAM 実行ロールを取得し、SageMaker セッションを作成して、ロール ARN とリージョンを出力します。

2. TGI NeuronX コンテナイメージの取得

from sagemaker.huggingface import get_huggingface_llm_image_uri
llm_image = get_huggingface_llm_image_uri(
    "huggingface-neuronx",
    version="0.0.20"
)
print(f"llm image uri: {llm_image}")

執筆時点では最新の DLC バージョンがヘルパー経由で公開されていなかったため、直接 ECR URI を使用しています。

3. Inferentia2 用にキャッシュ済みモデルをコンパイルまたは取得

Inferentia2 は動的シェイプをサポートしていないため、シーケンス長とバッチサイズはコンパイル時に固定する必要があります。Hugging Face は neuron model cache として事前コンパイル済み設定（例：Mistral‑7B、Zephyr‑7B）を提供しています。必要な設定が無い場合は、Optimum CLI を使ってコンパイルできます：

optimum-cli export neuron -m HuggingFaceH4/zephyr-7b-beta \
    --batch_size 4 --sequence_length 2048 \
    --num_cores 2 --auto_cast_type bf16 ./zephyr-7b-beta-neuron

コンパイルされた成果物は aws-neuron/zephyr-7b-seqlen-2048-bs-4-cores-2 として Hub にプッシュされます。

4. エンドポイント設定の定義とデプロイ

TGI NeuronX コンテナの主要環境変数は以下です：

  • HF_MODEL_ID, HF_NUM_CORES, HF_BATCH_SIZE, HF_SEQUENCE_LENGTH, HF_AUTO_CAST_TYPE
  • MAX_BATCH_SIZE, MAX_INPUT_LENGTH, MAX_TOTAL_TOKENS
from sagemaker.huggingface import HuggingFaceModel
config = {
    "HF_MODEL_ID": "HuggingFaceH4/zephyr-7b-beta",
    "HF_NUM_CORES": "2",
    "HF_BATCH_SIZE": "4",
    "HF_SEQUENCE_LENGTH": "2048",
    "HF_AUTO_CAST_TYPE": "bf16",
    "MAX_BATCH_SIZE": "4",
    "MAX_INPUT_LENGTH": "1512",
    "MAX_TOTAL_TOKENS": "2048",
}
llm_model = HuggingFaceModel(role=role, image_uri=llm_image, env=config)
llm = llm_model.deploy(initial_instance_count=1, instance_type="ml.inf2.8xlarge", container_startup_health_check_timeout=1800)

デプロイには通常 10〜15 分かかります。

5. 推論実行と Zephyr 7B とのチャット

モデルはチャットテンプレートを使用します。トークナイザーの apply_chat_template メソッドは OpenAI 形式のメッセージ辞書を必要なプロンプト形式に変換します。

from transformers import AutoTokenizer
tokenizer = AutoTokenizer.from_pretrained("aws-neuron/zephyr-7b-seqlen-2048-bs-4-cores-2")
messages = [
    {"role": "system", "content": "You are the AWS expert"},
    {"role": "user", "content": "Can you tell me an interesting fact about AWS?"},
]
prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
payload = {
    "do_sample": True,
    "top_p": 0.6,
    "temperature": 0.9,
    "top_k": 50,
    "max_new_tokens": 256,
    "repetition_penalty": 1.03,
    "return_full_text": False,
    "stop": ["</s>"]
}
chat = llm.predict({"inputs": prompt, "parameters": payload})
print(chat[0]["generated_text"][len(prompt):])

この応答により、モデルが Inferentia2 上で正常に動作していることが確認できます。

6. リソースのクリーンアップ

llm.delete_model()
llm.delete_endpoint()

今後の影響と課題

  • コスト効果の高いスケーリング: Inferentia2 は GPU クラスタに比べて、ハイスループット LLM 提供のためのより安価なコンピュートオプションを提供します。
  • モデルカバレッジ: 現在のキャッシュには Llama、Mistral、Zephyr が含まれており、Hugging Face はサポート対象アーキテクチャを拡大する予定です。
  • コンパイルワークフロー: ユーザーは事前キャッシュされたビルドを利用するか、Optimum を使用してカスタム設定をコンパイル（最大約45分）できます。
  • ロードマップ: 現在進行中の取り組みは、より多くのモデルのサポート、キャッシュシステムの改善、コンパイルパイプラインの効率化を目指しています。

ご質問は、Twitter または LinkedIn で Philipp Schmid にお問い合わせください。

Sources