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:

  1. Sigue la guía Configura el entorno de Cloud TPU para asegurarte de tener el acceso adecuado para usar Cloud TPU.

  2. Crea una identidad de servicio para la VM de TPU.

    gcloud alpha compute tpus tpu-vm service-identity create --zone=zone
  3. 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:

    1. Ve a la página Cuentas de servicio en la consola de Google Cloud.
    2. Haga clic en Crear cuenta de servicio.
    3. Ingresa el nombre de la cuenta de servicio.
    4. Opcional: Ingresa una descripción para la cuenta de servicio.
    5. Haz clic en Crear y continúa.
    6. Elige los roles que deseas otorgar a la cuenta de servicio.
    7. Haz clic en Continuar.
    8. (Opcional) Especifica los usuarios o grupos que pueden administrar la cuenta de servicio.
    9. 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 rol roles/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:

    1. Selecciona tu cuenta de servicio y haz clic en Agregar principal.
    2. En el campo Principales nuevas, ingresa la dirección de correo electrónico de tu cuenta de servicio.
    3. En el menú desplegable Seleccionar un rol, busca el rol (por ejemplo, Administrador de almacenamiento) y selecciónalo.
    4. Haz clic en Guardar.
  4. 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

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

  1. Crea un directorio para Miniconda:

    mkdir -p ~/miniconda3
  2. 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
  3. Instala Miniconda:

    bash ~/miniconda3/miniconda.sh -b -u -p ~/miniconda3
  4. Quita la secuencia de comandos del instalador de Miniconda:

    rm -rf ~/miniconda3/miniconda.sh
  5. Agrega Miniconda a tu variable PATH:

    export PATH="$HOME/miniconda3/bin:$PATH"
  6. Vuelve a cargar ~/.bashrc para aplicar los cambios a la variable PATH:

    source ~/.bashrc
  7. Crea un nuevo entorno de Conda:

    conda create -n tpu python=3.10
  8. Activa el entorno de Conda:

    source activate tpu

Cómo configurar MaxDiffusion

  1. Clona el repositorio de MaxDiffusion y navega al directorio MaxDiffusion:

    git clone https://github.com/google/maxdiffusion.git && cd maxdiffusion
  2. Cambia a la rama mlperf-4.1:

    git checkout mlperf4.1
  3. Instala MaxDiffusion:

    pip install -e .
  4. Instala las dependencias:

    pip install -r requirements.txt
  5. 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
  6. Instala dependencias adicionales:

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

Generar imágenes

  1. 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"
  2. 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