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:
Sigue la guía Configura el entorno de Cloud TPU para asegurarte de tener el acceso adecuado para usar Cloud TPU.
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. Puedes crear una cuenta de servicio desde la consola de Google Cloud o con el comando
gcloud
.Crea una cuenta de servicio con la herramienta de línea de comandos de
gcloud
:gcloud iam service-accounts create your-service-account-name \ --description="your-sa-description" \ --display-name="your-sa-display-name" export SERVICE_ACCOUNT_NAME=your-service-account-name
Crea una cuenta de servicio desde la consola de Google Cloud:
- Ve a la página Cuentas de servicio en la consola de Google Cloud.
- Haga clic en Crear cuenta de servicio.
- Ingresa el nombre de la cuenta de servicio.
- Opcional: Ingresa una descripción para la cuenta de servicio.
- Haz clic en Crear y continúa.
- Elige los roles que deseas otorgar a la cuenta de servicio.
- Haz clic en Continuar.
- (Opcional) Especifica los usuarios o grupos que pueden administrar la cuenta de servicio.
- Haz clic en Listo para terminar de crear la cuenta de servicio.
Después de crear tu cuenta de servicio, sigue estos pasos para otorgarle roles.
Se requieren los siguientes roles:
- Administrador de TPU: Es necesario para crear una TPU.
- Administrador de almacenamiento: Es necesario para acceder a Cloud Storage.
- Escritor de registros
- Escritor de métricas de Monitoring: Es necesario para escribir métricas en Cloud Monitoring.
Tu administrador debe otorgarte el rol
roles/resourcemanager.projectIamAdmin
para que puedas asignar roles de IAM a los usuarios. Un usuario con el rolroles/resourcemanager.projectIamAdmin
de administrador de IAM del proyecto también puede otorgar este rol.Usa los siguientes comandos
gcloud
para agregar roles de cuenta de servicio:gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/tpu.admin gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/storage.admin gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/logging.logWriter gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/monitoring.metricWriter
También puedes asignar roles con la consola de Google Cloud.
En la consola de Google Cloud, selecciona los siguientes roles:
- Selecciona tu cuenta de servicio y haz clic en Agregar principal.
- En el campo Principales nuevas, ingresa la dirección de correo electrónico de tu cuenta de servicio.
- En el menú desplegable Seleccionar un rol, busca el rol (por ejemplo, Administrador de almacenamiento) y selecciónalo.
- Haz clic en Guardar.
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
Cómo proteger la capacidad
Cuando tengas todo listo para proteger la capacidad de las TPU, revisa la página de cuotas para obtener información sobre el sistema de Cloud Quotas. Si tienes más preguntas sobre cómo reservar 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
- Este instructivo se probó con Python 3.10 o versiones posteriores.
- 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 . - 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 variablePATH
:source ~/.bashrc
Crea un nuevo entorno de Conda:
conda create -n tpu python=3.10
Activa el entorno de Conda:
source activate tpu
Cómo configurar 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