Optimum TPU を活用し、GKE 上で TPU を使用してオープンソース モデルをサービングします

このチュートリアルでは、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 は次の機能をサポートしています。

• 連続的なバッチ処理
• トークンのストリーミング
• トランスフォーマーを使用した貪欲探索と多項サンプリング。

モデルへのアクセス権を取得する

Gemma 2B モデルまたは Llama3 8B モデルを使用できます。このチュートリアルでは、この 2 つのモデルに焦点を当てていますが、Optimum TPU では他のモデルもサポートされています。

Gemma 2B

GKE にデプロイするために Gemma モデルへのアクセス権を取得するには、まずライセンス同意契約に署名してから、Hugging Face のアクセス トークンを生成する必要があります。

ライセンス同意契約に署名する

Gemma を使用するには同意契約に署名する必要があります。手順は次のとおりです。

1. モデルの同意ページにアクセスします。
2. Hugging Face アカウントを使用して同意を確認します。
3. モデルの規約に同意します。

アクセス トークンを生成する

Hugging Face トークンをまだ生成していない場合は、新しいトークンを生成します。

1. [Your Profile] > [Settings] > [Access Tokens] の順にクリックします。
2. [New Token] をクリックします。
3. 任意の名前と、少なくとも Read ロールを指定します。
4. [Generate a token] をクリックします。
5. トークンをクリップボードにコピーします。

Llama3 8B

Hugging Face リポジトリの Llama3 8b を使用するには、同意契約に署名する必要があります。

アクセス トークンを生成する

Hugging Face トークンをまだ生成していない場合は、新しいトークンを生成します。

1. [Your Profile] > [Settings] > [Access Tokens] の順にクリックします。
2. [New Token] を選択します。
3. 任意の名前と、少なくとも Read ロールを指定します。
4. [Generate a token] を選択します。
5. トークンをクリップボードにコピーします。

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

TPU リソースが使用可能な場合、GKE はノードプールをプロビジョニングします。TPU リソースが一時的に使用できない場合、出力には GCE_STOCKOUT エラー メッセージが表示されます。TPU ストックアウト エラーのトラブルシューティングについては、TPU リクエストに対応できる十分な TPU リソースがないをご覧ください。

クラスタと通信するように kubectl を構成する:

gcloud container clusters get-credentials ${CLUSTER_NAME} --location=${ZONE}

コンテナをビルドする

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 をデプロイする

このチュートリアルでは、Kubernetes Deployment を使用して Optimum TPU をデプロイします。Deployment は、クラスタ内のノードに分散された Pod の複数のレプリカを実行できる Kubernetes API オブジェクトです。

Gemma 2B

1. 次の Deployment マニフェストを optimum-tpu-gemma-2b-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=google/gemma-2b
           - --max-concurrent-requests=4
           - --max-input-length=8191
           - --max-total-tokens=8192
           - --max-batch-prefill-tokens=32768
           - --max-batch-size=16
           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 のデプロイを記述しています。

2. マニフェストを適用します。

   kubectl apply -f optimum-tpu-gemma-2b-2x4.yaml
   

Llama3 8B

1. 次のマニフェストを 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=8191
           - --max-total-tokens=8192
           - --max-batch-prefill-tokens=32768
           - --max-batch-size=16
           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 のデプロイを記述しています。

2. マニフェストを適用します。

   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"}