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

Inferencia de MaxDiffusion en TPU 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

Prepara el aprovisionamiento de una TPU v6e con 4 chips:

Accede a tu Cuenta de Google. Si aún no lo hiciste, regístrate para obtener una nueva cuenta.
En la consola de Google Cloud, selecciona o crea un proyecto Google Cloud en la página del selector de proyectos.
Habilita la facturación para tu Google Cloud proyecto. La facturación es obligatoria para todo el uso de Google Cloud .
Instala los componentes de gcloud alpha.
Ejecuta el siguiente comando para instalar la versión más reciente de los componentes de gcloud.
```
gcloud components update
```
Habilita la API de TPU con el siguiente comando de gcloud en Cloud Shell. También puedes habilitarlo desde la consola de Google Cloud.
```
gcloud services enable tpu.googleapis.com
```

Crea una identidad de servicio para la VM de TPU.

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

Crea una cuenta de servicio de TPU y otorga acceso a los servicios de Google Cloud .

Las cuentas de servicio permiten que el servicio de Google Cloud TPU acceda a otros servicios Google Cloud. Se recomienda una cuenta de servicio administrada por el usuario. Sigue estas guías para crear y otorgar roles. Se requieren los siguientes roles:
- Administrador de TPU: Necesario para crear una TPU
- Administrador de almacenamiento: Es necesario para acceder a Cloud Storage.
- Logs Writer: Es necesario para escribir registros con la API de Logging.
- Monitoring Metric Writer: Es necesario para escribir métricas en Cloud Monitoring.

Autentica 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

Cómo proteger la capacidad

Comunícate con tu equipo de ventas o de cuentas de Cloud TPU para solicitar la cuota de TPU y hacer preguntas sobre la capacidad.

Aprovisiona el entorno de Cloud TPU

Puedes aprovisionar TPUs v6e 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 en tu proyecto deGoogle Cloud .
Este instructivo se probó con la siguiente configuración:
- Python 3.10 or later
- Versiones de software nocturnas:
  - JAX nocturno 0.4.32.dev20240912
  - LibTPU nocturna 0.1.dev20240912+nightly
- Versiones de software estables:
  - JAX + JAX Lib de v0.4.35
Verifica que tu proyecto tenga suficiente cuota de TPU para lo siguiente:
- Cuota de VM de TPU
- Cuota de direcciones IP
- Quota de Hyperdisk Balanced
Permisos del proyecto del usuario
- Si usas GKE con XPK, consulta Permisos de la consola de Google Cloud 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 de tu 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.

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

Configura MaxDiffusion

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

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 entorno de ejecución de 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 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, asegúrate de limpiar los recursos de TPU.

Limpia

Borra la TPU:

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