Se usó la API de Cloud Translation para traducir esta página.

Inferencia de MaxDiffusion en TPUs v6e

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

Antes de comenzar

Prepárate para aprovisionar una TPU v6e con 4 chips:

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

Protege la capacidad

Cuando esté todo listo para proteger la capacidad de la TPU, consulta Cuotas de Cloud TPU para obtener más información sobre las cuotas de Cloud TPU. Si tienes más preguntas sobre cómo asegurar la capacidad, comunícate con tu equipo de ventas o de cuentas de Cloud TPU.

Aprovisiona el entorno de Cloud TPU

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

Requisitos previos

Verifica que tu proyecto tenga suficiente cuota de TPUS_PER_TPU_FAMILY, que especifica la cantidad máxima de chips a los que puedes acceder dentro de tu proyecto deGoogle Cloud .
Verifica que tu proyecto tenga suficiente cuota de TPU para lo siguiente:
- Cuota de VM de TPU
- Cuota de direcciones IP
- Cuota de Hyperdisk Balanced
Permisos del proyecto del usuario
- Si usas GKE con XPK, consulta Permisos de Cloud Console en la cuenta de usuario o de servicio para conocer los permisos necesarios para ejecutar XPK.

Aprovisiona 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 obtener una lista completa de los estados de las solicitudes de recursos en cola, consulta la documentación de Recursos en cola.

Conéctate a la TPU con SSH

   gcloud compute tpus tpu-vm ssh TPU_NAME

Crea 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

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

Agrega 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 nuevo:
```
conda create -n tpu python=3.10
```
Activa el entorno de Conda:
```
source activate tpu
```

Configura MaxDiffusion

Clona el repositorio de GitHub de MaxDiffusion y navega 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

Establece 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 instrucción y los parámetros de configuración definidos 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 generen las imágenes, asegúrate de limpiar los recursos de TPU.

Realiza una limpieza

Borra la TPU:

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