v6e TPU VM での JetStream PyTorch 推論

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

始める前に

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

  1. Cloud TPU 環境を設定するガイドに沿って、プロジェクトを設定し、Google Cloud CLI を構成し、Cloud TPU API を有効にして、Cloud TPU を使用するためのアクセス権があることを確認します。 Google Cloud

  2. 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 の割り当て
  • ユーザー プロジェクトの権限

環境変数を作成する

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 プロジェクト名。既存のプロジェクトを使用するか、新しいプロジェクトを作成します。
TPU_NAME TPU の名前。
ZONE サポートされているゾーンについては、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