このページは Cloud Translation API によって翻訳されました。

v6e TPU VM での JetStream MaxText 推論

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

始める前に

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

Google アカウントにログインします。Google アカウントをまだお持ちでない場合は、新しいアカウントを登録します。
Google Cloud コンソールで、プロジェクトセレクタページから Google Cloud プロジェクトを選択するか作成します。
Google Cloud プロジェクトに対する課金を有効にします。Google Cloud の使用にはすべて課金が必要です。
gcloud alpha components をインストールします。
次のコマンドを実行して、gcloud コンポーネントの最新バージョンをインストールします。
```
gcloud components update
```
Cloud Shell で、次の gcloud コマンドを使用して TPU API を有効にします。Google Cloud コンソールから有効にすることもできます。
```
gcloud services enable tpu.googleapis.com
```

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

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

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

サービスアカウントにより、Google Cloud TPU サービスが他の Google Cloud サービスにアクセスできるようになります。ユーザー管理のサービスアカウントの使用をおすすめします。ガイドに沿ってロールを作成し、付与します。次のロールが必要です。
- TPU 管理者: TPU を作成するために必要です
- ストレージ管理者: Cloud Storage にアクセスするために必要です
- ログ書き込み: Logging API でログを書き込むために必要です
- モニタリング指標の書き込み: Cloud Monitoring に指標を書き込むために必要
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 環境をプロビジョニングする

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

前提条件

プロジェクトに十分な TPUS_PER_TPU_FAMILY 割り当てがあることを確認します。これは、Google Cloud プロジェクト内でアクセスできるチップの最大数を指定します。
このチュートリアルは、次の構成でテストされました。
- Python 3.10 or later
- ナイトリーソフトウェアのバージョン:
  - 夜間 JAX 0.4.32.dev20240912
  - ナイトリー LibTPU 0.1.dev20240912+nightly
- 安定版ソフトウェアのバージョン:
  - v0.4.35 の JAX + JAX Lib
プロジェクトに次の TPU 割り当てがあることを確認します。
- TPU VM の割り当て
- IP アドレスの割り当て
- Hyperdisk Balanced の割り当て
ユーザープロジェクトの権限
- XPK で GKE を使用している場合は、XPK の実行に必要な権限について、ユーザーまたはサービスアカウントに対する Cloud コンソールの権限をご覧ください。

環境変数を作成する

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

TPU に接続したら、推論ベンチマークを実行できます。

Llama2-7B 推論ベンチマークを実行する

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

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

クリーンアップ

TPU を削除します。

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