v6e TPU VM での JetStream PyTorch 推論

このチュートリアルでは、JetStream を使用して TPU v6e で PyTorch モデルをサービングする方法について説明します。JetStream は、XLA デバイス（TPU）での大規模言語モデル（LLM）推論用にスループットとメモリが最適化されたエンジンです。このチュートリアルでは、Llama2-7B モデルの推論ベンチマークを実行します。

始める前に

4 個のチップを搭載した TPU v6e をプロビジョニングする準備を行います。

Cloud TPU 環境を設定するガイドに沿って、 Google Cloud プロジェクトを設定して、Google Cloud CLI の構成を行い、Cloud TPU API を有効にして、Cloud TPU を使用するためのアクセス権があることを確認します。
Google Cloud で認証し、Google Cloud CLI のデフォルトのプロジェクトとゾーンを構成します。
```
gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE
```

容量を確保する

TPU 容量を確保する準備ができたら、Cloud TPU 割り当てで Cloud TPU 割り当ての詳細を確認してください。容量の確保についてご不明な点がございましたら、Cloud TPU のセールスチームまたはアカウントチームにお問い合わせください。

Cloud TPU 環境をプロビジョニングする

TPU VM は、GKE、GKE と XPK、またはキューに格納されたリソースとしてプロビジョニングできます。

前提条件

プロジェクトに十分な TPUS_PER_TPU_FAMILY 割り当てがあることを確認します。これは、Google Cloud プロジェクト内でアクセスできるチップの最大数を指します。
プロジェクトの次の TPU 割り当てが十分にあることを確認します。
- TPU VM の割り当て
- IP アドレスの割り当て
- Hyperdisk Balanced の割り当て
ユーザープロジェクトの権限
- GKE と XPK を使用している場合は、XPK の実行に必要な権限について、ユーザーまたはサービスアカウントに対する Google Cloud コンソールの権限をご覧ください。

環境変数を作成する

Cloud Shell で、次の環境変数を作成します。

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-central2-b
export ACCELERATOR_TYPE=v6e-4
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

環境変数の説明

変数	説明
`PROJECT_ID`	実際の Google Cloud のプロジェクト ID。既存のプロジェクトを使用するか、新しいプロジェクトを作成します。
`TPU_NAME`	TPU の名前。
`ZONE`	TPU VM を作成するゾーン。サポートされているゾーンの詳細については、TPU のリージョンとゾーンをご覧ください。
`ACCELERATOR_TYPE`	アクセラレータタイプでは、作成する Cloud TPU のバージョンとサイズを指定します。TPU の各バージョンでサポートされているアクセラレータタイプの詳細については、TPU のバージョンをご覧ください。
`RUNTIME_VERSION`	Cloud TPU ソフトウェアのバージョン。
`SERVICE_ACCOUNT`	サービスアカウントのメールアドレス。これは、 Google Cloud コンソールの [サービスアカウント] ページで確認できます。例: `tpu-service-account@PROJECT_ID.iam.gserviceaccount.com`
`QUEUED_RESOURCE_ID`	キューに格納されたリソースリクエストのユーザー割り当てテキスト ID。

TPU v6e をプロビジョニングする

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
        --node-id ${TPU_NAME} \
        --project ${PROJECT_ID} \
        --zone ${ZONE} \
        --accelerator-type ${ACCELERATOR_TYPE} \
        --runtime-version ${RUNTIME_VERSION} \
        --service-account ${SERVICE_ACCOUNT}

list コマンドまたは describe コマンドを使用して、キューに格納されたリソースのステータスをクエリします。

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project ${PROJECT_ID} --zone ${ZONE}

キューに格納されたリソースリクエストのステータスの一覧については、キューに格納されたリソースのドキュメントをご覧ください。

SSH を使用して TPU に接続する

  gcloud compute tpus tpu-vm ssh ${TPU_NAME}

JetStream PyTorch Llama2-7B ベンチマークを実行する

JetStream-PyTorch を設定し、モデルチェックポイントを変換して推論ベンチマークを実行するには、GitHub リポジトリの手順に沿って操作します。

推論ベンチマークが完了したら、TPU リソースを必ずクリーンアップしてください。

クリーンアップ

TPU を削除します。

   gcloud compute tpus queued-resources delete ${QUEUED_RESOURCE_ID} \
      --project ${PROJECT_ID} \
      --zone ${ZONE} \
      --force \
      --async