Inferenza MaxDiffusion su TPU v6e

Questo tutorial mostra come pubblicare modelli MaxDiffusion su TPU v6e. In questo tutorial, generi immagini utilizzando il modello Stable Diffusion XL.

Prima di iniziare

Preparati a eseguire il provisioning di una TPU v6e con 4 chip:

  1. Accedi al tuo Account Google. Se non l'hai ancora fatto, registrati per creare un nuovo account.
  2. Nella console Google Cloud, seleziona o crea un Google Cloud progetto dalla pagina del selettore di progetti.
  3. Attiva la fatturazione per il tuo Google Cloud progetto. La fatturazione è obbligatoria per tutto l' Google Cloud utilizzo.
  4. Installa i componenti gcloud alpha.

  5. Esegui il seguente comando per installare la versione più recente dei componenti gcloud.

    gcloud components update

  6. Abilita l'API TPU tramite il seguente comando gcloud utilizzando Cloud Shell. Puoi anche attivarlo dalla console Google Cloud.

    gcloud services enable tpu.googleapis.com

  7. Crea un'identità di servizio per la VM TPU.

    gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE

  8. Crea un account di servizio TPU e concedi l'accesso ai Google Cloud servizi.

    Gli account di servizio consentono al Google Cloud servizio TPU di accedere ad altri Google Cloud servizi. È consigliato un account di servizio gestito dall'utente. Segui queste guide per creare e concedere i ruoli. Sono necessari i seguenti ruoli:

    • Amministratore TPU: necessario per creare una TPU
    • Amministratore dello spazio di archiviazione: necessario per accedere a Cloud Storage
    • Scrittore di log: necessario per scrivere i log con l'API Logging
    • Monitoring Metric Writer: necessario per scrivere le metriche in Cloud Monitoring

  9. Esegui l'autenticazione con Google Cloud e configura il progetto e la zona predefiniti per Google Cloud CLI.

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

Capacità sicura

Contatta il team di vendita o dedicato all'account Cloud TPU per richiedere una quota TPU e per eventuali domande sulla capacità.

Esegui il provisioning dell'ambiente Cloud TPU

Puoi eseguire il provisioning delle TPU v6e con GKE, con GKE e XPK, o come risorse in coda.

Prerequisiti

  • Verifica che il tuo progetto disponga di una quota TPUS_PER_TPU_FAMILY sufficiente, che specifica il numero massimo di chip a cui puoi accedere all'interno del progettoGoogle Cloud .
  • Questo tutorial è stato testato con la seguente configurazione:
    • Python 3.10 or later
    • Versioni software Nightly:
      • a notte JAX 0.4.32.dev20240912
      • nightly LibTPU 0.1.dev20240912+nightly
    • Versioni software stabili:
      • JAX + libreria JAX di v0.4.35
  • Verifica che il tuo progetto disponga di una quota TPU sufficiente per:
    • Quota VM TPU
    • Quota di indirizzi IP
    • Quota Hyperdisk bilanciata
  • Autorizzazioni del progetto per gli utenti

Provisiona una 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

Utilizza i comandi list o describe per eseguire query sullo stato della risorsa in coda.

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project ${PROJECT_ID} --zone ${ZONE}

Per un elenco completo degli stati delle richieste di risorse in coda, consulta la documentazione relativa alle risorse in coda.

Connettiti alla TPU tramite SSH

   gcloud compute tpus tpu-vm ssh TPU_NAME

Creare un ambiente Conda

  1. Crea una directory per Miniconda:

    mkdir -p ~/miniconda3

  2. Scarica lo script di installazione di Miniconda:

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

  3. Installa Miniconda:

    bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3

  4. Rimuovi lo script di installazione di Miniconda:

    rm -rf ~/miniconda3/miniconda.sh

  5. Aggiungi Miniconda alla variabile PATH:

    export PATH="$HOME/miniconda3/bin:$PATH"

  6. Ricarica ~/.bashrc per applicare le modifiche alla variabile PATH:

    source ~/.bashrc

  7. Crea un nuovo ambiente Conda:

    conda create -n tpu python=3.10

  8. Attiva l'ambiente Conda:

    source activate tpu

Configurare MaxDiffusion

  1. Clona il repository MaxDiffusion e vai alla directory MaxDiffusion:

    git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion

  2. Passa al branch mlperf-4.1:

    git checkout mlperf4.1

  3. Installa MaxDiffusion:

    pip install -e .

  4. Installa le dipendenze:

    pip install -r requirements.txt

  5. Installa 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. Installa le dipendenze aggiuntive:

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

Genera immagini

  1. Imposta le variabili di ambiente per configurare il 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. Genera le immagini utilizzando il prompt e le configurazioni definite in src/maxdiffusion/configs/base_xl.yml:

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

    Una volta generate le immagini, assicurati di ripulire le risorse TPU.

Esegui la pulizia

Elimina la TPU:

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