MaxDiffusion-Inferenz auf v6e-TPUs

In dieser Anleitung wird gezeigt, wie Sie MaxDiffusion-Modelle auf TPU v6e bereitstellen. In dieser Anleitung generieren Sie mit dem Modell „Stable Diffusion XL“ Bilder.

Vorbereitung

Bereiten Sie die Bereitstellung einer v6e-TPU mit 4 Chips vor:

Folgen Sie der Anleitung unter Cloud TPU-Umgebung einrichten, um ein Google Cloud -Projekt einzurichten, die Google Cloud CLI zu konfigurieren, die Cloud TPU API zu aktivieren und dafür zu sorgen, dass Sie Zugriff auf Cloud TPUs haben.
Authentifizieren Sie sich bei 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
```

Kapazität sichern

Wenn Sie sich TPU-Kapazität sichern möchten, finden Sie unter Cloud TPU-Kontingente weitere Informationen zu den Cloud TPU-Kontingenten. Wenn Sie weitere Fragen zur Sicherung von Kapazität 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 die Warteschlange gestellte Ressourcen bereitstellen.

Voraussetzungen

Prüfen Sie, ob Ihr Projekt über ein ausreichendes TPUS_PER_TPU_FAMILY-Kontingent verfügt. Damit wird die maximale Anzahl von Chips angegeben, auf die Sie in IhremGoogle Cloud -Projekt zugreifen können.
Prüfen Sie, ob folgende Kontingente für Ihr Projekt ausreichen:
- TPU-VM-Kontingent
- Kontingent für IP-Adressen
- Hyperdisk Balanced-Kontingent
Nutzerprojektberechtigungen
- Wenn Sie GKE mit XPK verwenden, lesen Sie den Artikel zu Cloud Console-Berechtigungen für das Nutzer- oder Dienstkonto, um zu erfahren, welche Berechtigungen 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 den Befehl list oder describe, um den Status Ihrer in die Warteschlange gestellten Ressource abzufragen.

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

Eine vollständige Liste der Status von Anfragen für in die Warteschlange gestellte Ressourcen finden Sie unter In die Warteschlange gestellte 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-Installations-Script 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-Installations-Script:
```
rm -rf ~/miniconda3/miniconda.sh
```
Fügen Sie Miniconda zu Ihrer PATH-Variablen hinzu:
```
export PATH="$HOME/miniconda3/bin:$PATH"
```
Aktualisieren Sie ~/.bashrc, um die Änderungen auf die Variable PATH 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-GitHub-Repository und wechseln Sie in das MaxDiffusion-Verzeichnis:
```
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
```
Wechseln Sie zum Zweig mlperf-4.1:
```
git checkout mlperf4.1
```
Installieren Sie MaxDiffusion:
```
pip install -e .
```
Installieren Sie die Abhängigkeiten:
```
pip install -r requirements.txt
```

Installieren Sie 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

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"

Generieren Sie mithilfe des Prompts und der in src/maxdiffusion/configs/base_xl.yml definierten Konfigurationen Bilder:
```
python -m src.maxdiffusion.generate_sdxl src/maxdiffusion/configs/base_xl.yml run_name="my_run"
```
Wenn 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