記事一覧に戻る

Philipp Schmidのプロフィール画像

David Corvoysier のプロフィール画像

Hugging Face Text Generation Inference (TGI) が AWS Inferentia2 と Amazon SageMaker で一般提供開始されたことをお知らせします。

Text Generation Inference (TGI) は、大規模言語モデル (LLM) を大規模な本番環境で展開・提供するための専用ソリューションです。TGI は Tensor Parallelism と continuous batching を活用し、Llama、Mistral など人気のオープン LLM 向けに高性能なテキスト生成を実現します。Text Generation Inference は、Grammarly、Uber、Deutsche Telekom など多くの企業で本番環境で利用されています。

Amazon SageMaker への TGI 統合と AWS Inferentia2 の組み合わせは、GPU に代わる強力なソリューションとなり、本番向け LLM アプリケーション構築のための実用的な選択肢となります。シームレスな統合によりモデルの展開と保守が容易になり、幅広い本番ユースケースで LLM をよりアクセスしやすく、スケーラブルにします。

AWS の新規 TGI for AWS Inferentia2 on Amazon SageMaker により、AWS のお客様は、HuggingChatOpenAssistant、および Hugging Face Hub 上の Serverless Endpoints for LLMs など、高い同時実行性と低遅延の LLM エクスペリエンスを支える技術と同じものを活用できます。

AWS Inferentia2 を使用した Amazon SageMaker への Zephyr 7B のデプロイ

本チュートリアルでは、Zephyr 7B などの最先端 LLM を AWS Inferentia2 と Amazon SageMaker を使用して簡単にデプロイする方法を紹介します。Zephyr は mistralai/Mistral-7B-v0.1 を基にした 7B パラメータのファインチューニング版で、Direct Preference Optimization (DPO) を用いて公開データセットと合成データセットの混合データで学習されています。詳細は 技術レポート をご覧ください。モデルは Apache 2.0 ライセンスで公開されており、幅広いアクセスと利用が可能です。

以下の手順を説明します:

  1. 開発環境のセットアップ
  2. TGI Neuronx イメージの取得
  3. Amazon SageMaker への Zephyr 7B のデプロイ
  4. 推論の実行とモデルとのチャット

それでは始めましょう。

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

Amazon SageMaker への Zephyr のデプロイには sagemaker Python SDK を使用します。AWS アカウントが設定されており、sagemaker Python SDK がインストールされていることを確認してください。

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

ローカル環境で SageMaker を使用する場合、SageMaker に必要な権限を持つ IAM ロールへのアクセスが必要です。詳細は こちら をご覧ください。

import sagemaker
import boto3
sess = sagemaker.Session()
# sagemaker session bucket -> used for uploading data, models and logs
# sagemaker will automatically create this bucket if it doesn't exist
sagemaker_session_bucket=None
if sagemaker_session_bucket is None and sess is not None:
    # set to default bucket if a bucket name is not given
    sagemaker_session_bucket = sess.default_bucket()

try:
    role = sagemaker.get_execution_role()
except ValueError:
    iam = boto3.client('iam')
    role = iam.get_role(RoleName='sagemaker_execution_role')['Role']['Arn']

sess = sagemaker.Session(default_bucket=sagemaker_session_bucket)

print(f"sagemaker role arn: {role}")
print(f"sagemaker session region: {sess.boto_region_name}")

2. TGI Neuronx イメージの取得

新しい Hugging Face TGI Neuronx DLC を使用して AWS Inferentia2 上で推論を実行できます。sagemaker SDK の get_huggingface_llm_image_uri メソッドを使用して、希望する backendsessionregionversion に基づいて適切な Hugging Face TGI Neuronx DLC URI を取得できます。利用可能な全バージョンは こちら で確認できます。

注:本ブログ投稿執筆時点では、最新版の Hugging Face LLM DLC は get_huggingface_llm_image_uri メソッドからまだ取得できません。代わりに生のコンテナ URI を使用します。

from sagemaker.huggingface import get_huggingface_llm_image_uri

# retrieve the llm image uri
llm_image = get_huggingface_llm_image_uri(
  "huggingface-neuronx",
  version="0.0.20"
)

# print ecr image uri
print(f"llm image uri: {llm_image}")

4. Amazon SageMaker への Zephyr 7B のデプロイ

Text Generation Inference (TGI) on Inferentia2 は Llama、Mistral など人気のオープン LLM をサポートしています。サポートされているモデルの完全な一覧(text-generation)は こちら で確認できます。

Inferentia2 向け LLM のコンパイル

執筆時点で、AWS Inferentia2 は推論時の動的シェイプをサポートしていません。そのため、シーケンス長とバッチサイズを事前に指定する必要があります。 Inferentia2 の全機能を活用しやすくするため、neuron model cache を作成しました。これは最も人気のある LLM の事前コンパイル済み設定を含んでいます。キャッシュされた設定は、モデルアーキテクチャ(Mistral)、モデルサイズ(7B)、neuron バージョン(2.16)、Inferentia コア数(2)、バッチサイズ(2)、シーケンス長(2048)で定義されます。

つまり、自分でモデルをコンパイルする必要はなく、キャッシュから事前コンパイル済みモデルを使用できます。例として mistralai/Mistral-7B-v0.1HuggingFaceH4/zephyr-7b-beta があります。コンパイル済み/キャッシュ済みの設定は Hugging Face Hub で確認できます。希望する設定がまだキャッシュされていない場合は、Optimum CLI を使用して自分でコンパイルするか、Cache リポジトリ でリクエストを開いてください。

本投稿では、inf2.8xlarge インスタンス上で以下のコマンドとパラメータを使用して HuggingFaceH4/zephyr-7b-beta を再コンパイルし、aws-neuron/zephyr-7b-seqlen-2048-bs-4-cores-2 にプッシュしました。

# compile model with optimum for batch size 4 and sequence length 2048
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
# push model to hub [repo_id] [local_path] [path_in_repo]
huggingface-cli upload  aws-neuron/zephyr-7b-seqlen-2048-bs-4 ./zephyr-7b-beta-neuron ./ --exclude "checkpoint/**"
# Move tokenizer to neuron model repository
python -c "from transformers import AutoTokenizer; AutoTokenizer.from_pretrained('HuggingFaceH4/zephyr-7b-beta').push_to_hub('aws-neuron/zephyr-7b-seqlen-2048-bs-4')"

まだキャッシュされていない設定で LLM をコンパイルする場合、最大 45 分かかることがあります。

TGI Neuronx エンドポイントのデプロイ

Amazon SageMaker へのモデルデプロイ前に、TGI Neuronx エンドポイント設定を定義する必要があります。以下の追加パラメータを定義してください:

  • HF_NUM_CORES: コンパイルに使用する Neuron コア数。
  • HF_BATCH_SIZE: モデルコンパイル時に使用したバッチサイズ。
  • HF_SEQUENCE_LENGTH: モデルコンパイル時に使用したシーケンス長。
  • HF_AUTO_CAST_TYPE: モデルコンパイル時に使用したオートキャストタイプ。

従来の TGI パラメータも定義する必要があります:

  • HF_MODEL_ID: Hugging Face モデル ID。
  • HF_TOKEN: ゲート付きモデルにアクセスするための Hugging Face API トークン。
  • MAX_BATCH_SIZE: モデルが扱える最大バッチサイズ(コンパイル時のバッチサイズと同一)。
  • MAX_INPUT_LENGTH: モデルが扱える最大入力長。
  • MAX_TOTAL_TOKENS: モデルが生成できる最大トークン数(コンパイル時に使用したシーケンス長と同一)。
import json
from sagemaker.huggingface import HuggingFaceModel

# sagemaker config & model config
instance_type = "ml.inf2.8xlarge"
health_check_timeout = 1800

# Define Model and Endpoint configuration parameter
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",
}

# create HuggingFaceModel with the image uri
llm_model = HuggingFaceModel(
  role=role,
  image_uri=llm_image,
  env=config
)

HuggingFaceModel を作成したら、deploy メソッドを使用して Amazon SageMaker にデプロイできます。ml.inf2.8xlarge インスタンスタイプでモデルをデプロイします。

# Deploy model to an endpoint
llm = llm_model.deploy(
  initial_instance_count=1,
  instance_type=instance_type,
  container_startup_health_check_timeout=health_check_timeout,
)

SageMaker はエンドポイントを作成し、そこにモデルをデプロイします。この処理には 10〜15 分かかることがあります。

5. 推論の実行とモデルとのチャット

エンドポイントのデプロイ完了後、predictorpredict メソッドを使用して推論を実行できます。生成に影響を与えるさまざまなパラメータをペイロードの parameters 属性に追加できます。サポートされているパラメータは こちら または TGI の OpenAPI 仕様(swagger ドキュメント)で確認できます。

HuggingFaceH4/zephyr-7b-beta は会話型チャットモデルです。以下のようなプロンプト構造でチャットできます:

<|system|>\nYou are a friendly.</s>\n<|user|>\nInstruction</s>\n<|assistant|>\n

手動でプロンプトを準備するのはエラーが発生しやすいため、トークナイザーの apply_chat_template メソッドを使用できます。これは OpenAI 形式のよく知られた messages 辞書を受け取り、モデルに適した形式に変換します。Zephyr が AWS に関する事実を知っているか確認してみましょう。

from transformers import AutoTokenizer

# load the tokenizer
tokenizer = AutoTokenizer.from_pretrained("aws-neuron/zephyr-7b-seqlen-2048-bs-4-cores-2")

# Prompt to generate
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)

# Generation arguments
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):])
# Sure, here's an interesting fact about AWS: As of 2021, AWS has more than 200 services in its portfolio, ranging from compute power and storage to databases,

素晴らしい、AWS Inferentia2 上の Amazon SageMaker に Zephyr を正常にデプロイし、チャットできました。

6. クリーンアップ

クリーンアップするには、モデルとエンドポイントを削除します。

llm.delete_model()
llm.delete_endpoint()

結論

Hugging Face Text Generation Inference (TGI) と AWS Inferentia2、Amazon SageMaker の統合により、大規模言語モデル (LLM) をデプロイするためのコスト効率の高い代替ソリューションが提供されます。

私たちは、より多くのモデルのサポート、コンパイルプロセスの簡素化、キャッシュシステムの改善に積極的に取り組んでいます。


ご一読ありがとうございます!ご質問があれば、Twitter または LinkedIn までお気軽にお問い合わせください。