Inferenza MaxDiffusion su TPU v6e

Attenzione: questa pagina si applica all'API Cloud TPU, che non è più in fase di sviluppo attivo. L'API Cloud TPU riceverà solo correzioni di bug e aggiornamenti della sicurezza. Per le funzionalità più recenti e il supporto per le versioni TPU più recenti, ti consigliamo di utilizzare Compute Engine o Google Kubernetes Engine per gestire le risorse TPU. Per saperne di più, consulta Risorse TPU in Compute Engine e Informazioni sulle TPU in GKE. L' API Cloud TPU non è supportata su TPU7x e versioni TPU successive.

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

Nota: dopo aver completato il benchmark di inferenza, assicurati di liberare spazio nelle risorse TPU.

Prima di iniziare

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

1. Segui Configurare l'ambiente Cloud TPU guida per configurare un Google Cloud progetto, configurare Google Cloud CLI, abilitare l'API Cloud TPU e assicurarti di avere accesso all'utilizzo delle Cloud TPU.

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

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

Garantisci la capacità

Quando sei pronto a garantire la capacità TPU, consulta Cloud TPU Quote per ulteriori informazioni sulle quote Cloud TPU. Se hai altre domande sulla protezione della capacità, contatta il team di vendita o il team dedicato all'account Cloud TPU.

Esegui il provisioning dell'ambiente Cloud TPU

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

Nota: questo documento descrive come eseguire il provisioning delle TPU utilizzando le risorse in coda. Se esegui il provisioning delle TPU utilizzando XPK (uno strumento CLI wrapper su GKE), configura le autorizzazioni XPK sul tuo account utente per GKE.

Prerequisiti

Nota: questo tutorial è stato testato con Python 3.10 o versioni successive.

• 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 Google Cloud progetto.
• Verifica che il tuo progetto disponga di una quota TPU sufficiente per:
  • Quota VM TPU
  • Quota indirizzi IP
  • Quota Hyperdisk bilanciato
• Autorizzazioni del progetto utente
  • Se utilizzi GKE con XPK, consulta Autorizzazioni di Cloud Console sull' account utente o di servizio per le autorizzazioni necessarie per eseguire XPK.

Esegui il provisioning di 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 una 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 utilizzando SSH

   gcloud compute tpus tpu-vm ssh TPU_NAME

Crea 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

Configura MaxDiffusion

1. Clona il repository GitHub 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 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

6. Installa 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 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 eseguire la pulizia delle risorse TPU.

Libera spazio

Elimina la TPU:

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