このチュートリアルでは、Hugging Face の Optimum TPU サービング フレームワークを活用し、Google Kubernetes Engine（GKE）上で Tensor Processing Unit（TPU）を使用して大規模言語モデル（LLM）オープンソース モデルをサービングする方法について説明します。このチュートリアルでは、Hugging Face からオープンソース モデルをダウンロードし、Optimum TPU を実行するコンテナを使用して、GKE Standard クラスタにモデルをデプロイします。
このガイドは、AI / ML ワークロードをデプロイしてサービングする際に、マネージド Kubernetes での詳細な制御、スケーラビリティ、復元力、ポータビリティ、費用対効果が求められる場合の出発点として適しています。
このチュートリアルは、Hugging Face エコシステムの生成 AI をご利用のお客様、GKE の新規または既存のユーザー、ML エンジニア、MLOps（DevOps）エンジニア、LLM のサービングに Kubernetes コンテナのオーケストレーション機能を使用することに関心をお持ちのプラットフォーム管理者を対象としています。
Google Cloud では、LLM 推論のためのオプションが複数用意されています。たとえば、Vertex AI、GKE、Google Compute Engine などのサービスを利用して、JetStream、vLLM などのサービング ライブラリやその他のパートナー サービスを組み込むことができます。たとえば、JetStream を使用してプロジェクトから最新の最適化を取得できます。Hugging Face のオプションを選択する場合は、Optimum TPU を使用できます。
Optimum TPU は次の機能をサポートしています。
- 連続的なバッチ処理
- トークンのストリーミング
- トランスフォーマーを使用した貪欲探索と多項サンプリング。
目標
- モデルの特性に基づいて推奨される TPU トポロジを持つ GKE Standard クラスタを準備します。
- GKE に Optimum TPU をデプロイします。
- Optimum TPU を使用して、サポートされるモデルを curl を通じてサービングします。
始める前に
- Sign in to your Google Cloud account. If you're new to Google Cloud, create an account to evaluate how our products perform in real-world scenarios. New customers also get $300 in free credits to run, test, and deploy workloads.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Make sure that billing is enabled for your Google Cloud project.
-
Enable the required API.
-
In the Google Cloud console, on the project selector page, select or create a Google Cloud project.
-
Make sure that billing is enabled for your Google Cloud project.
-
Enable the required API.
-
Make sure that you have the following role or roles on the project: roles/container.admin, roles/iam.serviceAccountAdmin
Check for the roles
-
In the Google Cloud console, go to the IAM page.
Go to IAM - Select the project.
-
In the Principal column, find all rows that identify you or a group that you're included in. To learn which groups you're included in, contact your administrator.
- For all rows that specify or include you, check the Role colunn to see whether the list of roles includes the required roles.
Grant the roles
-
In the Google Cloud console, go to the IAM page.
[IAM] に移動 - プロジェクトを選択します。
- [ アクセスを許可] をクリックします。
-
[新しいプリンシパル] フィールドに、ユーザー ID を入力します。 これは通常、Google アカウントのメールアドレスです。
- [ロールを選択] リストでロールを選択します。
- 追加のロールを付与するには、 [別のロールを追加] をクリックして各ロールを追加します。
- [保存] をクリックします。
-
- Hugging Face アカウントを作成します（まだ作成していない場合）。
- GKE の Cloud TPU 用にプロジェクトに十分な割り当てがあることを確認します。
環境を準備する
このチュートリアルでは、Cloud Shell を使用して Google Cloud でホストされるリソースを管理します。Cloud Shell には、このチュートリアルに必要な kubectl や gcloud CLI などのソフトウェアがプリインストールされています。
Cloud Shell を使用して環境を設定するには、次の操作を行います。
Google Cloud コンソールで、Google Cloud コンソールの
[Cloud Shell をアクティブにする] をクリックして、Cloud Shell セッションを起動します。これにより、Google Cloud コンソールの下部ペインでセッションが起動されます。デフォルトの環境変数を設定します。
gcloud config set project PROJECT_ID export PROJECT_ID=$(gcloud config get project) export CLUSTER_NAME=CLUSTER_NAME export REGION=REGION_NAME export ZONE=ZONE次の値を置き換えます。
- PROJECT_ID: Google Cloud のプロジェクト ID。
- CLUSTER_NAME: GKE クラスタの名前。
- REGION_NAME: GKE クラスタ、Cloud Storage バケット、TPU ノードが配置されているリージョン。TPU v5e マシンタイプを使用できるゾーン（
us-west1、us-west4、us-central1、us-east1、us-east5、europe-west4など）が含まれているリージョンです。 - （Standard クラスタのみ）ZONE: TPU リソースが使用可能なゾーン（
us-west4-aなど）。Autopilot クラスタの場合は、リージョンのみを指定します。ゾーンを指定する必要はありません。
Optimum TPU リポジトリのクローンを作成します。
git clone https://github.com/huggingface/optimum-tpu.git
モデルへのアクセス権を取得する
Gemma 2B モデルまたは Llama3 8B モデルを使用できます。このチュートリアルでは、この 2 つのモデルに焦点を当てていますが、Optimum TPU では他のモデルもサポートされています。
Gemma 2B
GKE にデプロイするために Gemma モデルへのアクセス権を取得するには、まずライセンス同意契約に署名してから、Hugging Face のアクセス トークンを生成する必要があります。
ライセンス同意契約に署名する
Gemma を使用するには同意契約に署名する必要があります。手順は次のとおりです。
- モデルの同意ページにアクセスします。
- Hugging Face アカウントを使用して同意を確認します。
- モデルの規約に同意します。
アクセス トークンを生成する
Hugging Face トークンをまだ生成していない場合は、新しいトークンを生成します。
- [Your Profile] > [Settings] > [Access Tokens] の順にクリックします。
- [New Token] をクリックします。
- 任意の名前と、少なくとも
Readロールを指定します。 - [Generate a token] をクリックします。
- トークンをクリップボードにコピーします。
Llama3 8B
Hugging Face リポジトリの Llama3 8b を使用するには、同意契約に署名する必要があります。
アクセス トークンを生成する
Hugging Face トークンをまだ生成していない場合は、新しいトークンを生成します。
- [Your Profile] > [Settings] > [Access Tokens] の順にクリックします。
- [New Token] を選択します。
- 任意の名前と、少なくとも
Readロールを指定します。 - [Generate a token] を選択します。
- トークンをクリップボードにコピーします。
GKE クラスタを作成する
1 つの CPU ノードを含む GKE Standard クラスタを作成します。
gcloud container clusters create CLUSTER_NAME \
--project=PROJECT_ID \
--num-nodes=1 \
--location=ZONE
TPU ノードプールを作成する
1 つのノードと 8 つのチップを含む v5e TPU ノードプールを作成します。
gcloud container node-pools create tpunodepool \
--location=ZONE \
--num-nodes=1 \
--machine-type=ct5lp-hightpu-8t \
--cluster=CLUSTER_NAME
クラスタと通信するように kubectl を設定する:
gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${REGION}
コンテナをビルドする
make コマンドを実行してイメージをビルドします。
cd optimum-tpu && make tpu-tgi
イメージを Artifact Registry に push する
gcloud artifacts repositories create optimum-tpu --repository-format=docker --location=REGION_NAME && \
gcloud auth configure-docker REGION_NAME-docker.pkg.dev && \
docker image tag huggingface/optimum-tpu REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest && \
docker push REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest
Hugging Face の認証情報用の Kubernetes Secret を作成する
Hugging Face トークンを含む Kubernetes Secret を作成します。
kubectl create secret generic hf-secret \
--from-literal=hf_api_token=${HF_TOKEN} \
--dry-run=client -o yaml | kubectl apply -f -
Optimum TPU をデプロイする
Optimum TPU をデプロイします。
Gemma 2B
次の Deployment マニフェストを
optimum-tpu-gemma-2b-2x4.yamlとして保存します。Deployment は、クラスタ内のノード間で Pod の複数のレプリカを実行できる Kubernetes API です。apiVersion: apps/v1 kind: Deployment metadata: name: tgi-tpu spec: replicas: 1 selector: matchLabels: app: tgi-tpu template: metadata: labels: app: tgi-tpu spec: nodeSelector: cloud.google.com/gke-tpu-topology: 2x4 cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice containers: - name: tgi-tpu image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest args: - --model-id=google/gemma-2b - --max-concurrent-requests=4 - --max-input-length=32 - --max-total-tokens=64 - --max-batch-size=1 securityContext: privileged: true env: - name: HF_TOKEN valueFrom: secretKeyRef: name: hf-secret key: hf_api_token ports: - containerPort: 80 resources: limits: google.com/tpu: 8 livenessProbe: httpGet: path: /health port: 80 initialDelaySeconds: 300 periodSeconds: 120 --- apiVersion: v1 kind: Service metadata: name: service spec: selector: app: tgi-tpu ports: - name: http protocol: TCP port: 8080 targetPort: 80このマニフェストは、TCP ポート 8080 に内部ロードバランサがある Optimum TPU のデプロイを記述しています。
マニフェストを適用します。
kubectl apply -f optimum-tpu-gemma-2b-2x4.yaml
Llama3 8B
次のマニフェストを
optimum-tpu-llama3-8b-2x4.yamlとして保存します。apiVersion: apps/v1 kind: Deployment metadata: name: tgi-tpu spec: replicas: 1 selector: matchLabels: app: tgi-tpu template: metadata: labels: app: tgi-tpu spec: nodeSelector: cloud.google.com/gke-tpu-topology: 2x4 cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice containers: - name: tgi-tpu image: REGION_NAME-docker.pkg.dev/PROJECT_ID/optimum-tpu/tgi-tpu:latest args: - --model-id=meta-llama/Meta-Llama-3-8B - --max-concurrent-requests=4 - --max-input-length=32 - --max-total-tokens=64 - --max-batch-size=1 env: - name: HF_TOKEN valueFrom: secretKeyRef: name: hf-secret key: hf_api_token ports: - containerPort: 80 resources: limits: google.com/tpu: 8 livenessProbe: httpGet: path: /health port: 80 initialDelaySeconds: 300 periodSeconds: 120 --- apiVersion: v1 kind: Service metadata: name: service spec: selector: app: tgi-tpu ports: - name: http protocol: TCP port: 8080 targetPort: 80このマニフェストは、TCP ポート 8080 に内部ロードバランサがある Optimum TPU のデプロイを記述しています。
マニフェストを適用します。
kubectl apply -f optimum-tpu-llama3-8b-2x4.yaml
実行中の Deployment のログを表示します。
kubectl logs -f -l app=tgi-tpu
出力例を以下に示します。
2024-07-09T22:39:34.365472Z WARN text_generation_router: router/src/main.rs:295: no pipeline tag found for model google/gemma-2b
2024-07-09T22:40:47.851405Z INFO text_generation_router: router/src/main.rs:314: Warming up model
2024-07-09T22:40:54.559269Z INFO text_generation_router: router/src/main.rs:351: Setting max batch total tokens to 64
2024-07-09T22:40:54.559291Z INFO text_generation_router: router/src/main.rs:352: Connected
2024-07-09T22:40:54.559295Z WARN text_generation_router: router/src/main.rs:366: Invalid hostname, defaulting to 0.0.0.0
モデルが完全にダウンロードされたことを確認してから、次のセクションに進んでください。
モデルをサービングする
モデルへのポート転送を設定します。
kubectl port-forward svc/service 8080:8080
curl を使用してモデルサーバーと対話する
デプロイされたモデルを確認します。
新しいターミナル セッションで curl を使用してモデルとチャットします。
curl 127.0.0.1:8080/generate -X POST -d '{"inputs":"What is Deep Learning?","parameters":{"max_new_tokens":40}}' -H 'Content-Type: application/json'
出力例を以下に示します。
{"generated_text":"\n\nDeep learning is a subset of machine learning that uses artificial neural networks to learn from data.\n\nArtificial neural networks are inspired by the way the human brain works. They are made up of multiple layers"}
クリーンアップ
このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。
デプロイされたリソースを削除する
このガイドで作成したリソースについて、Google Cloud アカウントに課金されないようにするには、次のコマンドを実行します。
gcloud container clusters delete CLUSTER_NAME \
--location=ZONE
次のステップ
- Optimum TPU のドキュメントを確認する。
- GKE で Gemma モデルを実行する方法と、GKE プラットフォームのオーケストレーション機能を使用して最適化された AI / ML ワークロードを実行する方法を確認する。
- GKE の TPU の詳細を確認する。