MaxDiffusion-Inferenz auf TPUs der Version 6e
In dieser Anleitung wird beschrieben, wie Sie MaxDiffusion-Modelle auf TPU v6e bereitstellen. In dieser Anleitung generieren Sie Bilder mit dem Stable Diffusion XL-Modell.
Hinweise
TPU v6e mit 4 Chips vorbereiten:
Folgen Sie der Anleitung Cloud TPU-Umgebung einrichten, um ein Google Cloud -Projekt einzurichten, die Google Cloud CLI zu konfigurieren, die Cloud TPU API zu aktivieren und sicherzustellen, dass Sie Zugriff auf Cloud TPUs haben.
Authentifizieren Sie sich mit Google Cloud und konfigurieren Sie das Standardprojekt und die Standardzone für die Google Cloud CLI.
gcloud auth login gcloud config set project PROJECT_ID gcloud config set compute/zone ZONE
Sichere Kapazität
Wenn Sie TPU-Kapazitäten reservieren möchten, finden Sie unter Cloud TPU-Kontingente weitere Informationen zu den Cloud TPU-Kontingenten. Wenn Sie weitere Fragen zur Kapazitätssicherung haben, wenden Sie sich an Ihr Cloud TPU-Vertriebs- oder Account-Management-Team.
Cloud TPU-Umgebung bereitstellen
Sie können TPU-VMs mit GKE, mit GKE und XPK oder als in der Warteschlange befindliche Ressourcen bereitstellen.
Vorbereitung
- Prüfen Sie, ob Ihr Projekt ein ausreichendes
TPUS_PER_TPU_FAMILY
-Kontingent hat. Dieses gibt die maximale Anzahl von Chips an, auf die Sie in IhremGoogle Cloud -Projekt zugreifen können. - Prüfen Sie, ob Ihr Projekt ein ausreichendes TPU-Kontingent für Folgendes hat:
- TPU-VM-Kontingent
- Kontingent für IP-Adressen
- Hyperdisk Balanced-Kontingent
- Nutzerberechtigungen für Projekte
- Wenn Sie GKE mit XPK verwenden, finden Sie unter Cloud Console-Berechtigungen für das Nutzer- oder Dienstkonto Informationen zu den Berechtigungen, die zum Ausführen von XPK erforderlich sind.
TPU v6e bereitstellen
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
Verwenden Sie die Befehle list
oder describe
, um den Status der in der Warteschlange befindlichen Ressource abzufragen.
gcloud alpha compute tpus queued-resources describe QUEUED_RESOURCE_ID \ --project=PROJECT_ID --zone=ZONE
Eine vollständige Liste der Status von anstehenden Ressourcenanfragen finden Sie in der Dokumentation zu anstehenden Ressourcen.
Über SSH eine Verbindung zur TPU herstellen
gcloud compute tpus tpu-vm ssh TPU_NAME
Conda-Umgebung erstellen
Erstellen Sie ein Verzeichnis für Miniconda:
mkdir -p ~/miniconda3
Laden Sie das Miniconda-Installationsskript herunter:
wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda3/miniconda.sh
Installieren Sie Miniconda:
bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
Entfernen Sie das Miniconda-Installationsskript:
rm -rf ~/miniconda3/miniconda.sh
Fügen Sie Miniconda der Variablen
PATH
hinzu:export PATH="$HOME/miniconda3/bin:$PATH"
Laden Sie
~/.bashrc
neu, um die Änderungen auf die VariablePATH
anzuwenden:source ~/.bashrc
Erstellen Sie eine neue Conda-Umgebung:
conda create -n tpu python=3.10
Aktivieren Sie die Conda-Umgebung:
source activate tpu
MaxDiffusion einrichten
Klonen Sie das MaxDiffusion-Repository und wechseln Sie zum MaxDiffusion-Verzeichnis:
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
Wechseln Sie zum
mlperf-4.1
-Zweig:git checkout mlperf4.1
Installieren Sie MaxDiffusion:
pip install -e .
Installieren Sie die Abhängigkeiten:
pip install -r requirements.txt
JAX installieren:
pip install jax[tpu]==0.4.34 jaxlib==0.4.34 ml-dtypes==0.2.0 -f https://storage.googleapis.com/jax-releases/jax_nightly_releases.html -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Installieren Sie zusätzliche Abhängigkeiten:
pip install huggingface_hub==0.25 absl-py flax tensorboardX google-cloud-storage torch tensorflow transformers
Bilder erstellen
Legen Sie Umgebungsvariablen fest, um die TPU-Laufzeit zu konfigurieren:
LIBTPU_INIT_ARGS="--xla_tpu_rwb_fusion=false --xla_tpu_dot_dot_fusion_duplicated=true --xla_tpu_scoped_vmem_limit_kib=65536"
Bilder mit dem Prompt und den Konfigurationen generieren, die in
src/maxdiffusion/configs/base_xl.yml
definiert sind:python -m src.maxdiffusion.generate_sdxl src/maxdiffusion/configs/base_xl.yml run_name="my_run"
Nachdem die Bilder generiert wurden, müssen Sie die TPU-Ressourcen bereinigen.
Bereinigen
Löschen Sie die TPU:
gcloud compute tpus queued-resources delete QUEUED_RESOURCE_ID \ --project PROJECT_ID \ --zone ZONE \ --force \ --async