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

v6e TPU での MaxDiffusion 推論

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

始める前に

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 コンソールの権限をご覧ください。

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
```
PATH 変数に Miniconda を追加します。
```
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