v6e TPU VM での JetStream PyTorch 推論

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

始める前に

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

  1. Cloud TPU 環境を設定するガイドに沿って、Cloud TPU を使用するための適切なアクセス権があることを確認します。

  2. TPU VM のサービス ID を作成します。

    gcloud alpha compute tpus tpu-vm service-identity create --zone=zone

  3. TPU サービス アカウントを作成し、 Google Cloud サービスへのアクセス権を付与します。

    サービス アカウントにより、 Google Cloud TPU サービスが他の Google Cloudサービスにアクセスできるようになります。ユーザー管理のサービス アカウントの使用をおすすめします。サービス アカウントは、Google Cloud コンソールまたは gcloud コマンドを使用して作成できます。

    gcloud コマンドライン ツールを使用してサービス アカウントを作成します。

    gcloud iam service-accounts create your-service-account-name \
--description="your-sa-description" \
--display-name="your-sa-display-name"
export SERVICE_ACCOUNT_NAME=your-service-account-name

    Google Cloud コンソールからサービス アカウントを作成します。

    1. Google Cloud コンソールで [サービス アカウント] ページに移動します。
    2. [サービス アカウントを作成] をクリックします。
    3. サービス アカウント名を入力します。
    4. （省略可）サービス アカウントの説明を入力します。
    5. [作成] をクリックして続行します。
    6. サービス アカウントに付与するロールを選択します。
    7. [続行] をクリックします。
    8. （省略可）サービス アカウントを管理できるユーザーまたはグループを指定します。
    9. [完了] をクリックして、サービス アカウントの作成を完了します。

    サービス アカウントを作成したら、次の手順でサービス アカウントのロールを付与します。

    次のロールが必要です。

    • TPU 管理者: TPU を作成するために必要です
    • ストレージ管理者: Cloud Storage にアクセスするために必要です。
    • ログ書き込み
    • モニタリング指標の書き込み: Cloud Monitoring に指標を書き込むために必要

    ユーザーに IAM ロールを割り当てるには、管理者から roles/resourcemanager.projectIamAdmin が付与されている必要があります。プロジェクト IAM 管理者 roles/resourcemanager.projectIamAdmin ロールを持つユーザーも、このロールを付与できます。

    次の gcloud コマンドを使用して、サービス アカウントのロールを追加します。

    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
   --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
   --role roles/tpu.admin
gcloud projects add-iam-policy-binding ${PROJECT_ID} \
   --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
   --role roles/storage.admin
gcloud projects add-iam-policy-binding ${PROJECT_ID} \
   --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
   --role roles/logging.logWriter
gcloud projects add-iam-policy-binding ${PROJECT_ID} \
   --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
   --role roles/monitoring.metricWriter

    Google Cloud コンソールを使用してロールを割り当てることもできます。

    Google Cloud コンソールで、次のロールを選択します。

    1. サービス アカウントを選択し、[プリンシパルを追加] をクリックします。
    2. [新しいプリンシパル] フィールドに、サービス アカウントのメールアドレスを入力します。
    3. [ロールを選択] プルダウンで、ロール（ストレージ管理者など）を検索して選択します。
    4. [保存] をクリックします。

  4. Google Cloud で認証し、Google Cloud CLI のデフォルトのプロジェクトとゾーンを構成します。

    gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE

容量を確保する

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

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

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

前提条件

  • このチュートリアルは Python 3.10 以降でテストされています。
  • プロジェクトに十分な TPUS_PER_TPU_FAMILY 割り当てがあることを確認します。これは、Google Cloud プロジェクト内でアクセスできるチップの最大数を指定します。
  • プロジェクトに次の TPU 割り当てがあることを確認します。
    • TPU VM の割り当て
    • IP アドレスの割り当て
    • Hyperdisk Balanced の割り当て
  • ユーザー プロジェクトの権限

環境変数を作成する

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

export NODE_ID=TPU_NODE_ID # TPU name
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=v6e-4
export ZONE=us-central2-b
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=YOUR_SERVICE_ACCOUNT
export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID
export VALID_DURATION=VALID_DURATION

# Additional environment variable needed for Multislice:
export NUM_SLICES=NUM_SLICES

# Use a custom network for better performance as well as to avoid having the
# default network becoming overloaded.
export NETWORK_NAME=${PROJECT_ID}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw

コマンドフラグの説明

変数 説明
NODE_ID キューに入れられたリソース リクエストの割り当て時に作成される TPU のユーザー割り当て ID。
PROJECT_ID Google Cloud プロジェクト名。既存のプロジェクトを使用するか、新しいプロジェクトを作成します。
ゾーン サポートされているゾーンについては、TPU のリージョンとゾーンのドキュメントをご覧ください。
ACCELERATOR_TYPE サポートされているアクセラレータ タイプについては、アクセラレータの種類のドキュメントをご覧ください。
RUNTIME_VERSION v2-alpha-tpuv6e
SERVICE_ACCOUNT これは、Google Cloud コンソール -> IAM -> サービス アカウント で確認できるサービス アカウントのメールアドレスです。
例: tpu-service-account@<your_project_ID>.iam.gserviceaccount.com.com
NUM_SLICES 作成するスライス数（マルチスライスの場合のみ必要）
QUEUED_RESOURCE_ID キューに格納されたリソース リクエストのユーザー割り当てテキスト ID。
VALID_DURATION キューに入れられたリソース リクエストが有効である期間。
NETWORK_NAME 使用するセカンダリ ネットワークの名前。
NETWORK_FW_NAME 使用するセカンダリ ネットワーク ファイアウォールの名前。

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

    gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
        --node-id TPU_NAME \
        --project PROJECT_ID \
        --zone ZONE \
        --accelerator-type v6e-4 \
        --runtime-version v2-alpha-tpuv6e \
        --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