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

v6e TPU での MaxDiffusion 推論

このチュートリアルでは、TPU v6e で MaxDiffusion モデルを提供する方法について説明します。このチュートリアルでは、Stable Diffusion XL モデルを使用して画像を生成します。

始める前に

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

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

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

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

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. [保存] をクリックします。
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 の割り当て
ユーザープロジェクトの権限
- XPK で GKE を使用している場合は、XPK の実行に必要な権限について、ユーザーまたはサービスアカウントに対する Cloud コンソールの権限をご覧ください。

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

Conda 環境を作成する

Miniconda のディレクトリを作成します。
```
mkdir -p ~/miniconda3
```

Miniconda インストーラスクリプトをダウンロードします。

wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh

Miniconda をインストールします。

bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3

Miniconda インストーラスクリプトを削除します。
```
rm -rf ~/miniconda3/miniconda.sh
```
Miniconda を PATH 変数に追加します。
```
export PATH="$HOME/miniconda3/bin:$PATH"
```
~/.bashrc を再読み込みして、PATH 変数に変更を適用します。
```
source ~/.bashrc
```
新しい Conda 環境を作成します。
```
conda create -n tpu python=3.10
```
Conda 環境を有効にします。
```
source activate tpu
```

MaxDiffusion を設定する

MaxDiffusion リポジトリのクローンを作成し、MaxDiffusion ディレクトリに移動します。
```
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
```
mlperf-4.1 ブランチに切り替えます。
```
git checkout mlperf4.1
```
MaxDiffusion をインストールします。
```
pip install -e .
```
依存関係をインストールします。
```
pip install -r requirements.txt
```

JAX をインストールします。

pip install -U --pre jax[tpu] -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

追加の依存関係をインストールします。

 pip install huggingface_hub==0.25 absl-py flax tensorboardX google-cloud-storage torch tensorflow transformers

画像を生成

環境変数を設定して TPU ランタイムを構成します。

LIBTPU_INIT_ARGS="--xla_tpu_rwb_fusion=false --xla_tpu_dot_dot_fusion_duplicated=true --xla_tpu_scoped_vmem_limit_kib=65536"

src/maxdiffusion/configs/base_xl.yml で定義されたプロンプトと構成を使用して画像を生成します。
```
python -m src.maxdiffusion.generate_sdxl src/maxdiffusion/configs/base_xl.yml run_name="my_run"
```
画像が生成されたら、TPU リソースを必ずクリーンアップしてください。

クリーンアップ

TPU を削除します。

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