Cloud TPU で Hex-LLM プレミアムコンテナを使用してオープンモデルを提供する

XLA を使用して高効率の大規模言語モデル（LLM）のサービスを提供する Hex-LLM は、Cloud TPU ハードウェア用に設計および最適化された Vertex AI LLM サービングフレームワークです。Hex-LLM は、継続的なバッチ処理や PagedAttention などの LLM サービングテクノロジーと、XLA と Cloud TPU 向けに調整された Vertex AI の最適化を組み合わせています。これは、オープンソースモデル用の Cloud TPU で、効率的で低コストの LLM サービスを提供します。

Hex-LLM は、モデルプレイグラウンド、ワンクリックデプロイ、ノートブックを通じて Model Garden で利用できます。

機能

Hex-LLM は、XLA と Cloud TPU 向けの Google 独自の最適化を備えたオープンソースプロジェクトに基づいています。Hex-LLM は、頻繁に使用される LLM のサービスを提供する際に、高いスループットと低レイテンシを実現します。

Hex-LLM には次の最適化が含まれています。

多数の同時リクエストでモデルがハードウェアを十分活用できるようにするトークンベースの連続バッチ処理アルゴリズム。
XLA 用に最適化されたアテンションカーネルの完全な置き換え。
複数の Cloud TPU チップで LLM を効率的に実行するために、高度に最適化された重みシャーディング手法による柔軟でコンポーザブルなデータ並列処理とテンソル並列処理の戦略。

Hex-LLM は、高密度の LLM からスパース LLM まで幅広くサポートしています。

Gemma 2B および 7B
Gemma 2 9B および 27B
Llama 2 7B、13B、70B
Llama 3 8B および 70B
Mistral 7B および Mixtral 8x7B

Hex-LLM には、次のようなさまざまな機能も用意されています。

Hex-LLM は単一のコンテナに含まれています。Hex-LLM は、API サーバー、推論エンジン、サポートされているモデルを単一の Docker イメージにパッケージ化してデプロイします。
Hugging Face モデル形式に対応しています。Hex-LLM は、ローカルディスク、Hugging Face Hub、Cloud Storage バケットから Hugging Face モデルを読み込むことができます。
bitsandbytes と AWQ を使用した量子化。
動的 LoRA 読み込み。Hex-LLM は、サービング中にリクエスト引数を読み取ることで LoRA 重みを読み込むことができます。

Model Garden を使ってみる

Hex-LLM Cloud TPU サービングコンテナは Model Garden に統合されています。このサービングテクノロジーは、さまざまなモデルのプレイグラウンド、ワンクリックデプロイ、Colab Enterprise ノートブックの例で利用できます。

プレイグラウンドを使用する

Model Garden のプレイグラウンドは、事前にデプロイされた Vertex AI エンドポイントであり、モデルカードでリクエストを送信することでアクセスできます。

プロンプトを入力し、必要に応じてリクエストの引数を指定します。
[送信] をクリックして、モデルのレスポンスをすばやく取得します。

Gemma で試す

ワンクリックデプロイを使用する

モデルカードを使用して、Hex-LLM を備えたカスタム Vertex AI エンドポイントをデプロイできます。

モデルカードページに移動し、[デプロイ] をクリックします。
使用するモデルのバリエーションに対して、デプロイするための Cloud TPU v5e マシンタイプを選択します。
下部にある [デプロイ] をクリックして、デプロイプロセスを開始します。2 通のメール通知が届きます。1 通はモデルがアップロードされたとき、もう 1 通はエンドポイントの準備が整ったときです。

Colab Enterprise ノートブックを使用する

柔軟性とカスタマイズのために、Colab Enterprise ノートブックの例を使用して、Vertex AI SDK for Python を使用して Hex-LLM で Vertex AI エンドポイントをデプロイできます。

モデルカードのページに移動し、[ノートブックを開く] をクリックします。
Vertex Serving ノートブックを選択します。ノートブックが Colab Enterprise で開きます。
ノートブックを実行して、Hex-LLM を使用してモデルをデプロイし、予測リクエストをエンドポイントに送信します。デプロイのコードスニペットは次のとおりです。

hexllm_args = [
    f"--model=google/gemma-2-9b-it",
    f"--tensor_parallel_size=4",
    f"--hbm_utilization_factor=0.8",
    f"--max_running_seqs=512",
]
hexllm_envs = {
    "PJRT_DEVICE": "TPU",
    "MODEL_ID": "google/gemma-2-9b-it",
    "DEPLOY_SOURCE": "notebook",
}
model = aiplatform.Model.upload(
    display_name="gemma-2-9b-it",
    serving_container_image_uri=HEXLLM_DOCKER_URI,
    serving_container_command=[
        "python", "-m", "hex_llm.server.api_server"
    ],
    serving_container_args=hexllm_args,
    serving_container_ports=[7080],
    serving_container_predict_route="/generate",
    serving_container_health_route="/ping",
    serving_container_environment_variables=hexllm_envs,
    serving_container_shared_memory_size_mb=(16 * 1024),
    serving_container_deployment_timeout=7200,
)

endpoint = aiplatform.Endpoint.create(display_name="gemma-2-9b-it-endpoint")
model.deploy(
    endpoint=endpoint,
    machine_type="ct5lp-hightpu-4t",
    deploy_request_timeout=1800,
    service_account="<your-service-account>",
    min_replica_count=1,
    max_replica_count=1,
)

カスタムサービングでは、次の Hex-LLM サーバー起動引数を変更できます。

--model: 読み込むモデル。Hugging Face モデル ID、ローカルの絶対パス、Cloud Storage バケットパスを指定できます。
--tokenizer: 読み込むトークナイザ。Hugging Face モデル ID、ローカルの絶対パス、Cloud Storage バケットパスを指定できます。デフォルト値は --model と同じです。
--enable_jit: JIT モードを有効にするかどうか。デフォルト値は True です。
--enable_lora: LoRA 読み込みモードを有効にするかどうか。デフォルト値は False です。
--max_lora_rank: LoRA アダプターでサポートされている最大 LoRA ランク。リクエストで定義されます。デフォルト値は 16 です。
--data_parallel_size: データ並列レプリカの数。デフォルト値は 1 です。
--tensor_parallel_size: テンソル並列レプリカの数。デフォルト値は 1 です。
--max_running_seqs: サーバーが同時に処理できるリクエストの最大数。この引数を大きくすると、サーバーが達成できるスループットが高くなりますが、レイテンシに悪影響を及ぼす可能性があります。デフォルト値は 256 です。
--hbm_utilization_factor: モデルの重みが読み込まれた後に KV キャッシュに割り当てることができる無料の Cloud TPU HBM の割合。この引数を小さい値に設定すると、Cloud TPU HBM のメモリ不足を効果的に防ぐことができます。デフォルト値は 0.9 です。
--seed: すべての乱数生成ツールを初期化するためのシード。この引数を変更すると、同じプロンプトで生成される出力に影響する可能性があります。デフォルト値は 0 です。

Cloud TPU の割り当てをリクエストする

Model Garden では、デフォルトの割り当ては us-west1 リージョンの Cloud TPU v5e チップ 4 つです。この割り当ては、ワンクリックデプロイと Colab Enterprise ノートブックのデプロイに適用されます。割り当ての追加をリクエストするには、割り当ての増加をリクエストするをご覧ください。

Cloud TPU で Hex-LLM プレミアム コンテナを使用してオープンモデルを提供する

機能