Inferencia de JetStream PyTorch en VMs de TPU v6e

En este instructivo, se muestra cómo usar JetStream para entregar modelos de PyTorch en TPU v6e. JetStream es un motor con capacidad de procesamiento y memoria optimizada para la inferencia de modelos de lenguaje grandes (LLM) en dispositivos XLA (TPU). En este instructivo, ejecutarás la comparativa de inferencia para el modelo Llama2-7B.

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
  • 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.

Crea variables de entorno

En Cloud Shell, crea las siguientes variables de entorno:

export NODE_ID=TPU_NODE_ID # TPU name
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=v6e-4
export ZONE=us-central2-b
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=YOUR_SERVICE_ACCOUNT
export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID
export VALID_DURATION=VALID_DURATION

# Additional environment variable needed for Multislice:
export NUM_SLICES=NUM_SLICES

# Use a custom network for better performance as well as to avoid having the
# default network becoming overloaded.
export NETWORK_NAME=${PROJECT_ID}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw

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

Ejecuta la comparativa Llama2-7B de JetStream PyTorch

Para configurar JetStream-PyTorch, convertir los puntos de control del modelo y ejecutar la comparativa de inferencia, sigue las instrucciones en el repositorio de GitHub.

Cuando se complete la comparativa de inferencia, 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