Esta página se ha traducido con Cloud Translation API.

Inferencia de MaxDiffusion en TPUs v6e

En este tutorial se muestra cómo servir modelos de MaxDiffusion en TPU v6e. En este tutorial, generarás imágenes con el modelo Stable Diffusion XL.

Antes de empezar

Prepara el aprovisionamiento de una TPU v6e con 4 chips:

Sigue la guía Configurar el entorno de TPU de Cloud para configurar un proyecto de Google Cloud , configurar la CLI de Google Cloud, habilitar la API Cloud TPU y asegurarte de que tienes acceso para usar las TPU de Cloud.
Autentícate con Google Cloud y configura el proyecto y la zona predeterminados de la CLI de Google Cloud.
```
gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE
```

Capacidad segura

Cuando quieras proteger la capacidad de las TPUs, consulta Cuotas de TPUs de Cloud para obtener más información sobre las cuotas de TPUs de Cloud. Si tienes más preguntas sobre cómo proteger la capacidad, ponte en contacto con el equipo de Ventas o de Cuentas de Cloud TPU.

Aprovisionar el entorno de TPU de Cloud

Puedes aprovisionar VMs de TPU con GKE, con GKE y XPK o como recursos en cola.

Requisitos previos

Comprueba que tu proyecto tenga suficiente cuota de TPUS_PER_TPU_FAMILY, que especifica el número máximo de chips a los que puedes acceder en tuGoogle Cloud proyecto.
Comprueba que tu proyecto tenga suficiente cuota de TPU para lo siguiente:
- Cuota de máquinas virtuales de TPU
- Cuota de direcciones IP
- Cuota de Hyperdisk Balanced
Permisos de proyecto de usuario
- Si usas GKE con XPK, consulta los permisos de la consola de Cloud en la cuenta de usuario o de servicio para saber qué permisos necesitas para ejecutar XPK.

Aprovisionar 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

Usa los comandos list o describe para consultar el estado del recurso en cola.

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

Para ver una lista completa de los estados de las solicitudes de recursos en cola, consulta la documentación sobre recursos en cola.

Conectarse a la TPU mediante SSH

   gcloud compute tpus tpu-vm ssh TPU_NAME

Crear un entorno de Conda

Crea un directorio para Miniconda:
```
mkdir -p ~/miniconda3
```

Descarga la secuencia de comandos del instalador de Miniconda:

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

Instala Miniconda:

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

Elimina la secuencia de comandos del instalador de Miniconda:
```
rm -rf ~/miniconda3/miniconda.sh
```

Añade Miniconda a tu variable PATH:

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

Vuelve a cargar ~/.bashrc para aplicar los cambios a la variable PATH:
```
source ~/.bashrc
```
Crea un entorno de Conda:
```
conda create -n tpu python=3.10
```
Activa el entorno de Conda:
```
source activate tpu
```

Configurar MaxDiffusion

Clona el repositorio de GitHub de MaxDiffusion y ve al directorio de MaxDiffusion:
```
git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
```
Cambia a la rama mlperf-4.1:
```
git checkout mlperf4.1
```
Instala MaxDiffusion:
```
pip install -e .
```
Instala las dependencias:
```
pip install -r requirements.txt
```

Instala 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

Instala dependencias adicionales:

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

Generar imágenes

Define variables de entorno para configurar el tiempo de ejecución de la TPU:

LIBTPU_INIT_ARGS="--xla_tpu_rwb_fusion=false --xla_tpu_dot_dot_fusion_duplicated=true --xla_tpu_scoped_vmem_limit_kib=65536"

Genera imágenes con la petición y las configuraciones definidas en src/maxdiffusion/configs/base_xl.yml:
```
python -m src.maxdiffusion.generate_sdxl src/maxdiffusion/configs/base_xl.yml run_name="my_run"
```
Cuando se hayan generado las imágenes, elimina los recursos de TPU.

Limpieza

Elimina la TPU:

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