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:

  1. Folgen Sie der Anleitung Cloud TPU-Umgebung einrichten, um sicherzustellen, dass Sie die erforderlichen Zugriffsrechte für die Verwendung von Cloud TPUs haben.

  2. Erstellen Sie eine Dienstidentität für die TPU-VM.

    gcloud alpha compute tpus tpu-vm service-identity create --zone=zone
  3. Erstellen Sie ein TPU-Dienstkonto und gewähren Sie Zugriff auf Google Cloud -Dienste.

    Dienstkonten ermöglichen dem Google Cloud TPU-Dienst, auf andere Google Cloud Dienste zuzugreifen. Ein nutzerverwaltetes Dienstkonto wird empfohlen. Sie können ein Dienstkonto in der Google Cloud Console oder über den Befehl gcloud erstellen.

    Erstellen Sie ein Dienstkonto mit dem gcloud-Befehlszeilentool:

    gcloud iam service-accounts create your-service-account-name \
    --description="your-sa-description" \
    --display-name="your-sa-display-name"
    export SERVICE_ACCOUNT_NAME=your-service-account-name

    So erstellen Sie ein Dienstkonto über die Google Cloud Console:

    1. Rufen Sie in der Google Cloud Console die Seite „Dienstkonten“ auf.
    2. Klicken Sie auf Dienstkonto erstellen.
    3. Geben Sie den Namen des Dienstkontos ein.
    4. Optional: Geben Sie eine Beschreibung für das Dienstkonto ein.
    5. Klicken Sie auf Erstellen und fahren Sie fort.
    6. Wählen Sie die Rollen aus, die Sie dem Dienstkonto zuweisen möchten.
    7. Klicken Sie auf Weiter.
    8. Optional: Geben Sie Nutzer oder Gruppen an, die das Dienstkonto verwalten können.
    9. Klicken Sie auf Fertig, um das Erstellen des Dienstkontos abzuschließen.

    Nachdem Sie das Dienstkonto erstellt haben, können Sie ihm Rollen gewähren.

    Folgende Rollen sind erforderlich:

    • TPU-Administrator: Erforderlich zum Erstellen einer TPU
    • Storage Admin: Erforderlich für den Zugriff auf Cloud Storage
    • Log-Autor
    • Monitoring-Messwert-Autor: Erforderlich zum Schreiben von Messwerten in Cloud Monitoring

    Ihr Administrator muss Ihnen die Rolle roles/resourcemanager.projectIamAdmin zuweisen, damit Sie Nutzern IAM-Rollen zuweisen können. Auch ein Nutzer mit der Rolle „Projekt-IAM-Administrator“ roles/resourcemanager.projectIamAdmin kann diese Rolle gewähren.

    Verwenden Sie die folgenden gcloud-Befehle, um Dienstkontorollen hinzuzufügen:

    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
       --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
       --role roles/tpu.admin
    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
       --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
       --role roles/storage.admin
    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
       --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
       --role roles/logging.logWriter
    gcloud projects add-iam-policy-binding ${PROJECT_ID} \
       --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \
       --role roles/monitoring.metricWriter

    Sie können Rollen auch über die Google Cloud Console zuweisen.

    Wählen Sie in der Google Cloud Console die folgenden Rollen aus:

    1. Wählen Sie Ihr Dienstkonto aus und klicken Sie auf Hauptkonto hinzufügen.
    2. Geben Sie im Feld Neue Hauptkonten die E-Mail-Adresse Ihres Dienstkontos ein.
    3. Suchen Sie im Drop-down-Menü Rolle auswählen nach einer Rolle (z. B. Storage-Administrator) und wählen Sie sie aus.
    4. Klicken Sie auf Speichern.
  4. 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 sichern möchten, lesen Sie den Artikel Kontingente, um mehr über das Cloud-Kontingentsystem zu erfahren. Wenn Sie weitere Fragen zur Kapazitätssicherung haben, wenden Sie sich an Ihr Cloud TPU-Vertriebs- oder Account-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

  • Diese Anleitung wurde mit Python 3.10 oder höher getestet.
  • 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

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

  1. Erstellen Sie ein Verzeichnis für Miniconda:

    mkdir -p ~/miniconda3
  2. Laden Sie das Miniconda-Installationsskript herunter:

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

    bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
  4. Entfernen Sie das Miniconda-Installationsskript:

    rm -rf ~/miniconda3/miniconda.sh
  5. Fügen Sie Miniconda der Variablen PATH hinzu:

    export PATH="$HOME/miniconda3/bin:$PATH"
  6. Laden Sie ~/.bashrc neu, um die Änderungen auf die Variable PATH anzuwenden:

    source ~/.bashrc
  7. Erstellen Sie eine neue Conda-Umgebung:

    conda create -n tpu python=3.10
  8. Aktivieren Sie die Conda-Umgebung:

    source activate tpu

MaxDiffusion einrichten

  1. Klonen Sie das MaxDiffusion-Repository und wechseln Sie zum MaxDiffusion-Verzeichnis:

    git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
  2. Wechseln Sie zum mlperf-4.1-Zweig:

    git checkout mlperf4.1
  3. Installieren Sie MaxDiffusion:

    pip install -e .
  4. Installieren Sie die Abhängigkeiten:

    pip install -r requirements.txt
  5. JAX installieren:

    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. Installieren Sie zusätzliche Abhängigkeiten:

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

Bilder erstellen

  1. 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"
  2. 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