TPU で vLLM を使用して Qwen2-7B-Instruct をサービングする

このチュートリアルでは、v6e TPU VM で vLLM TPU サービング フレームワークを使用して Qwen/Qwen2-7B-Instruct モデルをサービングします。

目標

1. 環境をセットアップする。
2. Qwen2-7B-Instruct で vLLM を実行します。
3. 推論リクエストを送信します。
4. ベンチマーク ワークロードを実行します。
5. クリーンアップする。

費用

このチュートリアルでは、Google Cloud の課金対象となる以下のコンポーネントを使用します。

• Cloud TPU
• Compute Engine

料金計算ツールを使うと、予想使用量に基づいて費用の見積もりを生成できます。

始める前に

このチュートリアルに進む前に、Cloud TPU 環境を設定するの手順に沿って操作してください。この手順では、 Google Cloudプロジェクトを作成し、Cloud TPU を使用するように構成するために必要な手順について説明します。既存のGoogle Cloud プロジェクトを使用することもできます。その場合は、Google Cloud プロジェクトの作成手順をスキップして、Cloud TPU を使用するための環境を設定するから開始できます。

このチュートリアルを使用するには、Hugging Face アクセス トークンが必要です。無料アカウントは Hugging Face で登録できます。アカウントを取得したら、アクセス トークンを生成します。

1. [Welcome to Hugging Face] ページで、アカウントのアバターをクリックして [アクセス トークン] を選択します。
2. [アクセス トークン] ページで、[新しいトークンを作成] をクリックします。
3. [読み取り] トークン タイプを選択し、トークンの名前を入力します。
4. アクセス トークンが表示されます。トークンは安全な場所に保存してください。

環境の設定

キューに格納されたリソース

キューに格納されたリソース API を使用して Cloud TPU v6e VM を作成します。qwen2-7b-instruct の場合は、v6e-1 TPU を使用することをおすすめします。

export PROJECT_ID=<PROJECT>
export TPU_NAME=<TPU_NAME>
export ZONE=<ZONE>
export QR_ID=<QR_ID>
export TPU_TYPE=<TPU_TYPE>

変数を設定します。

• PROJECT - プロジェクトの名前。
• TPU_NAME - 作成する TPU VM マシンの名前。
• ZONE - 新しい VM を作成するクラウド ゾーン。
• TPU_TYPE - 作成する TPU VM のタイプ。たとえば、v6e-1、v6e-4。
• QR_ID - 作成するキューに格納されたリソースの名前。

キューに格納されたリソース リクエストを作成します。

gcloud alpha compute tpus queued-resources create $QR_ID \
 --node-id $TPU_NAME \
 --project $PROJECT_ID \
 --zone $ZONE \
 --accelerator-type $TPU_TYPE \
 --runtime-version v2-alpha-tpuv6e

TPU VM の準備が整っていることを確認します。

gcloud compute tpus queued-resources describe $QR_ID \
  --project $PROJECT_ID \
  --zone $ZONE

ヒント: 最初は、キューに格納されたリソース リクエストのステータスは WAITING_FOR_RESOURCES になります。TPU VM が作成されると、ステータスが ACTIVE に変わります。リソースの可用性によっては、この処理に時間がかかることがあります。待ち時間が非常に長い場合は、別のゾーンを試すことを検討してください。ステータスが ACTIVE になるまで待ってから、続行します。

たとえば、ステータスが ACTIVE の場合:

name: projects/your-project-id/locations/your-zone/queuedResources/your-queued-resource-id
  state:
  state: ACTIVE
  tpu:
  nodeSpec:
  - node:
      acceleratorType: v6e-1
      bootDisk: {}
      networkConfig:
          enableExternalIps: true
      queuedResource: projects/your-project-number/locations/your-zone/queuedResources/your-queued-resource-id
      runtimeVersion: v2-alpha-tpuv6e
      schedulingConfig: {}
      serviceAccount: {}
      shieldedInstanceConfig: {}
      useTpuVm: true
      nodeId: your-node-id
      parent: projects/your-project-number/locations/your-zone

予約

予約を使用して Cloud TPU v6e VM を作成します。qwen2-7b-instruct の場合は、v6e-1 TPU を使用することをおすすめします。まず、環境変数を設定します。

export PROJECT_ID="YOUR_PROJECT_ID"
export TPU_NAME="TPU_MACHINE_NAME"
export ZONE="ZONE_NAME"
export TPU_TYPE=v6e-8
export HF_TOKEN="YOUR_HF_TOKEN"
export RESERVATION="RESERVATION_NAME"

変数を設定します。

• PROJECT - プロジェクトの名前。
• TPU_NAME - 作成する TPU VM マシンの名前。
• ZONE - 新しい VM を作成するクラウド ゾーン。
• TPU_TYPE - 作成する TPU VM のタイプ。たとえば、v6e-1、v6e-4。
• RESERVATION - TPU を含む予約の名前。

予約を使用して TPU VM を作成します。

gcloud alpha compute tpus tpu-vm create $TPU_NAME \
    --zone=$ZONE \
    --project $PROJECT_ID \
    --accelerator-type=$TPU_TYPE \
    --version=v2-alpha-tpuv6e \
    --provisioning-model=reservation-bound \
    --reservation=$RESERVATION

TPU VM に接続します。

gcloud compute tpus tpu-vm ssh $TPU_NAME \
  --project $PROJECT_ID \
  --zone $ZONE

Qwen2-7B-instruct で vLLM を実行する

1. Hugging Face トークンとモデル名の変数を設定します。

     export HF_TOKEN="YOUR_HF_TOKEN"
     export MODEL_NAME="Qwen/Qwen2-7B-Instruct"
   

2. TPU VM 内で、vLLM Docker コンテナをデタッチ モードで実行し、vLLM サーバーを起動します。このコマンドは、共有メモリサイズ 10 GB を使用します。

   export DOCKER_URI="vllm/vllm-tpu:v0.18.0"
   export CONTAINER_NAME="${USER}-vllm"
   export MAX_MODEL_LEN=4096
   export TP=1 # number of chips
   
   sudo docker run -d --name "${CONTAINER_NAME}" \
       --privileged --net=host \
       -v /dev/shm:/dev/shm \
       --shm-size 10gb \
       -e "HF_HOME=/dev/shm" \
       -e "HF_TOKEN=${HF_TOKEN}" \
       -p 8000:8000 "${DOCKER_URI}" \
           vllm serve ${MODEL_NAME} \
               --seed 42 \
               --gpu-memory-utilization 0.98 \
               --max-num-batched-tokens 1024 \
               --max-num-seqs 128 \
               --tensor-parallel-size $TP \
               --max-model-len $MAX_MODEL_LEN

3. サーバーログを確認して、実行されていることを確認します。

   sudo docker logs -f "${CONTAINER_NAME}"

   vLLM サーバーが実行されると、次のような出力が表示されます。出力が表示されたら、CTRL+C を押してターミナルに戻ります。

   (APIServer pid=7) INFO:     Started server process [7]
   (APIServer pid=7) INFO:     Waiting for application startup.
   (APIServer pid=7) INFO:     Application startup complete.
   

推論リクエストを送信する

vLLM サーバーが起動したら、API にリクエストを送信できます。詳細については、vLLM API リファレンス ドキュメントをご覧ください。

1. curl を使用して、サーバーにテスト リクエストを送信します。

   sudo docker exec "${CONTAINER_NAME}" \
     curl http://localhost:8000/v1/completions \
       -H "Content-Type: application/json" \
       -d '{
           "prompt": "The future of AI is",
           "max_tokens": 200,
           "temperature": 0
         }'

レスポンスは JSON 形式で返されます。

ベンチマーク ワークロードを実行する

2 つ目のターミナルから、実行中のサーバーに対してベンチマークを実行できます。

1. コンテナ内で datasets ライブラリをインストールします。

   sudo docker exec "${CONTAINER_NAME}" pip install datasets

2. コンテナ内で vllm bench serve コマンドを実行します。

   sudo docker exec "${CONTAINER_NAME}" \
       vllm bench serve \
           --backend vllm \
           --dataset-name random \
           --num-prompts 1000 \
           --seed 100

ベンチマーク結果は次のようになります。

============ Serving Benchmark Result ============
Successful requests:                     1000
Benchmark duration (s):                  45.35
Total input tokens:                      1024000
Total generated tokens:                  126848
Request throughput (req/s):              22.05
Output token throughput (tok/s):         2797.15
Peak output token throughput (tok/s):    4258.00
Peak concurrent requests:                1000.00
Total Token throughput (tok/s):          25377.57
---------------Time to First Token----------------
Mean TTFT (ms):                          21332.46
Median TTFT (ms):                        21330.37
P99 TTFT (ms):                           42436.47
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          37.36
Median TPOT (ms):                        38.56
P99 TPOT (ms):                           38.69
---------------Inter-token Latency----------------
Mean ITL (ms):                           37.35
Median ITL (ms):                         38.55
P99 ITL (ms):                            39.43
==================================================

クリーンアップ

このチュートリアルで使用したリソースについて、Google Cloud アカウントに課金されないようにするには、リソースを含むプロジェクトを削除するか、プロジェクトを維持して個々のリソースを削除します。

1. ターミナルで「exit」と入力して、TPU VM との接続を解除します。

リソースを削除する

プロジェクトを削除してすべてのリソースを削除することも、プロジェクトを保持してリソースを削除することもできます。

プロジェクトの削除

Google Cloud プロジェクトと関連するすべてのリソースを削除するには、次のコマンドを実行します。

    gcloud projects delete $PROJECT_ID

TPU リソースを削除する

キューに格納されたリソース

Cloud TPU リソースを削除します。次のコマンドは、--force パラメータを使用して、キューに格納されたリソース リクエストと TPU VM の両方を削除します。

gcloud alpha compute tpus queued-resources delete $QR_ID \
  --project=$PROJECT_ID \
  --zone=$ZONE \
  --force

予約

Cloud TPU VM を削除します。次のコマンドを使用して VM を終了し、TPU を予約に戻します。

gcloud compute tpus tpu-vm delete $TPU_NAME --zone $ZONE --project $PROJECT_ID --quiet

注: 予約がある限り、TPU の料金は引き続きお支払いいただきます。

次のステップ

• Cloud TPU 上の vLLM の詳細を確認する。
• Cloud TPU の詳細を確認する。