Inferensi MaxDiffusion di TPU v6e

Tutorial ini menunjukkan cara menayangkan model MaxDiffusion di TPU v6e. Dalam tutorial ini, Anda akan membuat gambar menggunakan model Stable Diffusion XL.

Sebelum memulai

Bersiap untuk menyediakan TPU v6e dengan 4 chip:

  1. Ikuti panduan Menyiapkan lingkungan Cloud TPU untuk menyiapkan Google Cloud project, mengonfigurasi Google Cloud CLI, mengaktifkan Cloud TPU API, dan memastikan Anda memiliki akses untuk menggunakan Cloud TPU.

  2. Lakukan autentikasi dengan Google Cloud dan konfigurasikan project dan zona default untuk Google Cloud CLI.

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

Kapasitas aman

Jika Anda siap untuk mengamankan kapasitas TPU, lihat Kuota Cloud TPU untuk mengetahui informasi selengkapnya tentang kuota Cloud TPU. Jika ada pertanyaan tambahan tentang cara mendapatkan kapasitas, hubungi tim akun atau tim penjualan Cloud TPU Anda.

Menyediakan lingkungan Cloud TPU

Anda dapat menyediakan VM TPU dengan GKE, dengan GKE dan XPK, atau sebagai resource dalam antrean.

Prasyarat

  • Pastikan project Anda memiliki kuota TPUS_PER_TPU_FAMILY yang cukup, yang menentukan jumlah maksimum chip yang dapat Anda akses dalam projectGoogle Cloud .
  • Pastikan project Anda memiliki cukup kuota TPU untuk:
    • Kuota VM TPU
    • Kuota alamat IP
    • Kuota Hyperdisk Balanced
  • Izin project pengguna

Menyediakan 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

Gunakan perintah list atau describe untuk membuat kueri status resource yang diantrekan.

   gcloud alpha compute tpus queued-resources describe QUEUED_RESOURCE_ID  \
      --project=PROJECT_ID --zone=ZONE

Untuk mengetahui daftar lengkap status permintaan resource yang diantrekan, lihat dokumentasi Resource yang Diantrekan.

Menghubungkan ke TPU menggunakan SSH

   gcloud compute tpus tpu-vm ssh TPU_NAME

Membuat lingkungan Conda

  1. Buat direktori untuk Miniconda:

    mkdir -p ~/miniconda3
  2. Download skrip penginstal Miniconda:

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

    bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
  4. Hapus skrip penginstal Miniconda:

    rm -rf ~/miniconda3/miniconda.sh
  5. Tambahkan Miniconda ke variabel PATH Anda:

    export PATH="$HOME/miniconda3/bin:$PATH"
  6. Muat ulang ~/.bashrc untuk menerapkan perubahan pada variabel PATH:

    source ~/.bashrc
  7. Buat lingkungan Conda baru:

    conda create -n tpu python=3.10
  8. Aktifkan lingkungan Conda:

    source activate tpu

Menyiapkan MaxDiffusion

  1. Clone repositori MaxDiffusion dan buka direktori MaxDiffusion:

    git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
  2. Beralih ke cabang mlperf-4.1:

    git checkout mlperf4.1
  3. Instal MaxDiffusion:

    pip install -e .
  4. Instal dependensi:

    pip install -r requirements.txt
  5. Menginstal 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
  6. Instal dependensi tambahan:

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

Buat gambar

  1. Tetapkan variabel lingkungan untuk mengonfigurasi runtime TPU:

    LIBTPU_INIT_ARGS="--xla_tpu_rwb_fusion=false --xla_tpu_dot_dot_fusion_duplicated=true --xla_tpu_scoped_vmem_limit_kib=65536"
  2. Buat gambar menggunakan perintah dan konfigurasi yang ditentukan di src/maxdiffusion/configs/base_xl.yml:

    python -m src.maxdiffusion.generate_sdxl src/maxdiffusion/configs/base_xl.yml run_name="my_run"

    Setelah gambar dibuat, pastikan untuk membersihkan resource TPU.

Pembersihan

Hapus TPU:

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