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:
- Login ke Akun Google Anda. Jika Anda belum melakukannya, daftar untuk membuat akun baru.
- Di konsol Google Cloud, pilih atau buat project Google Cloud dari halaman pemilih project.
- Aktifkan penagihan untuk project Google Cloud Anda. Penagihan diperlukan untuk semua penggunaan Google Cloud .
- Instal komponen gcloud alpha.
Jalankan perintah berikut untuk menginstal versi terbaru komponen
gcloud
.gcloud components update
Aktifkan TPU API melalui perintah
gcloud
berikut menggunakan Cloud Shell. Anda juga dapat mengaktifkannya dari konsol .gcloud services enable tpu.googleapis.com
Buat identitas layanan untuk VM TPU.
gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE
Buat akun layanan TPU dan berikan akses ke layanan Google Cloud .
Akun layanan memungkinkan layanan TPU Google Cloud untuk mengakses layanan Google Cloud lainnya. Akun layanan yang dikelola pengguna direkomendasikan. Ikuti panduan ini untuk membuat dan memberikan peran. Peran berikut diperlukan:
- TPU Admin: Diperlukan untuk membuat TPU
- Storage Admin: Diperlukan untuk mengakses Cloud Storage
- Logs Writer: Diperlukan untuk menulis log dengan Logging API
- Monitoring Metric Writer: Diperlukan untuk menulis metrik ke Cloud Monitoring
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
Hubungi tim penjualan atau akun Cloud TPU Anda untuk meminta kuota TPU dan mengajukan pertanyaan tentang kapasitas.
Menyediakan lingkungan Cloud TPU
Anda dapat menyediakan TPU v6e dengan GKE, dengan GKE dan XPK, atau sebagai resource dalam antrean.
Prasyarat
- Pastikan project Anda memiliki cukup kuota
TPUS_PER_TPU_FAMILY
, yang menentukan jumlah maksimum chip yang dapat Anda akses dalam projectGoogle Cloud . - Tutorial ini diuji dengan konfigurasi berikut:
- Python
3.10 or later
- Versi software harian:
- JAX harian
0.4.32.dev20240912
- LibTPU harian
0.1.dev20240912+nightly
- JAX harian
- Versi software stabil:
- JAX + JAX Lib dari
v0.4.35
- JAX + JAX Lib dari
- Python
- Pastikan project Anda memiliki cukup kuota TPU untuk:
- Kuota VM TPU
- Kuota Alamat IP
- Kuota Hyperdisk Balanced
- Izin project pengguna
- Jika Anda menggunakan GKE dengan XPK, lihat Izin Konsol Cloud di akun pengguna atau layanan untuk mengetahui izin yang diperlukan untuk menjalankan XPK.
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
Buat direktori untuk Miniconda:
mkdir -p ~/miniconda3
Download skrip penginstal Miniconda:
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
Instal Miniconda:
bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
Hapus skrip penginstal Miniconda:
rm -rf ~/miniconda3/miniconda.sh
Tambahkan Miniconda ke variabel
PATH
Anda:export PATH="$HOME/miniconda3/bin:$PATH"
Muat ulang
~/.bashrc
untuk menerapkan perubahan pada variabelPATH
:source ~/.bashrc
Buat lingkungan Conda baru:
conda create -n tpu python=3.10
Aktifkan lingkungan Conda:
source activate tpu
Menyiapkan MaxDiffusion
Clone repositori MaxDiffusion dan buka direktori MaxDiffusion:
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
Beralih ke cabang
mlperf-4.1
:git checkout mlperf4.1
Instal MaxDiffusion:
pip install -e .
Instal dependensi:
pip install -r requirements.txt
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
Instal dependensi tambahan:
pip install huggingface_hub==0.25 absl-py flax tensorboardX google-cloud-storage torch tensorflow transformers
Buat gambar
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"
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