在 v6e TPU 上執行 MaxDiffusion 推論

本教學課程將說明如何在 TPU v6e 上提供 MaxDiffusion 模型。在本教學課程中，您將使用 Stable Diffusion XL 模型產生圖片。

事前準備

準備佈建 4 晶片的 TPU v6e：

  1. 請按照「設定 Cloud TPU 環境」指南設定 Google Cloud 專案、設定 Google Cloud CLI、啟用 Cloud TPU API，並確保您有權使用 Cloud TPU。

  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 環境

您可以使用 GKE、GKE 和 XPK，或以佇列資源佈建 TPU VM。

必要條件

  • 請確認專案具備足夠的 TPUS_PER_TPU_FAMILY 配額，該配額會指定您在Google Cloud 專案中可存取的晶片數量上限。
  • 確認專案具備足夠的 TPU 配額，可供下列項目使用：
    • TPU VM 配額
    • IP 位址配額
    • Hyperdisk Balanced 配額
  • 使用者專案權限

配置 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 環境

  1. 建立 Miniconda 的目錄：

    mkdir -p ~/miniconda3

  2. 下載 Miniconda 安裝程式指令碼：

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

  3. 安裝 Miniconda：

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

  4. 移除 Miniconda 安裝程式指令碼：

    rm -rf ~/miniconda3/miniconda.sh

  5. 將 Miniconda 新增至 PATH 變數：

    export PATH="$HOME/miniconda3/bin:$PATH"

  6. 重新載入 ~/.bashrc，將變更套用至 PATH 變數：

    source ~/.bashrc

  7. 建立新的 Conda 環境：

    conda create -n tpu python=3.10

  8. 啟用 Conda 環境：

    source activate tpu

設定 MaxDiffusion

  1. 複製 MaxDiffusion GitHub 存放區，然後前往 MaxDiffusion 目錄：

    git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion

  2. 切換至 mlperf-4.1 分支版本：

    git checkout mlperf4.1

  3. 安裝 MaxDiffusion：

    pip install -e .

  4. 安裝依附元件：

    pip install -r requirements.txt

  5. 安裝 JAX：

    pip install jax[tpu]==0.4.34 jaxlib==0.4.34 ml-dtypes==0.2.0 -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ -f https://storage.googleapis.com/jax-releases/libtpu_releases.html

  6. 安裝其他依附元件：

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

生成圖像

  1. 設定環境變數以設定 TPU 執行階段：

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

  2. 使用 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