Inferencia de PyTorch de JetStream en VMs de TPU v6e

En este tutorial se muestra cómo usar JetStream para servir modelos de PyTorch en TPU v6e. JetStream es un motor optimizado de rendimiento y memoria para la inferencia de modelos de lenguaje extensos (LLMs) en dispositivos XLA (TPUs). En este tutorial, ejecutarás la prueba comparativa de inferencia del modelo Llama2-7B.

Nota: Después de completar la prueba de rendimiento de inferencia, no olvides limpiar los recursos de TPU.

Antes de empezar

Prepara el aprovisionamiento de una TPU v6e con 4 chips:

1. Sigue la guía Configurar el entorno de TPU de Cloud para configurar un proyecto de Google Cloud , configurar la CLI de Google Cloud, habilitar la API Cloud TPU y asegurarte de que tienes acceso para usar las TPU de Cloud.

2. Autentícate con Google Cloud y configura el proyecto y la zona predeterminados de la CLI de Google Cloud.

   gcloud auth login
   gcloud config set project PROJECT_ID
   gcloud config set compute/zone ZONE

Capacidad segura

Cuando quieras proteger la capacidad de las TPUs, consulta Cuotas de TPUs de Cloud para obtener más información sobre las cuotas de TPUs de Cloud. Si tienes más preguntas sobre cómo proteger la capacidad, ponte en contacto con el equipo de Ventas o de Cuentas de Cloud TPU.

Aprovisionar el entorno de TPU de Cloud

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

Nota: En este documento se describe cómo aprovisionar TPUs mediante recursos en cola. Si vas a aprovisionar tus TPUs mediante XPK (una herramienta de CLI de envoltorio sobre GKE), configura los permisos de XPK en tu cuenta de usuario para GKE.

Requisitos previos

Nota: Este tutorial se ha probado con Python 3.10 o versiones posteriores.

• Comprueba que tu proyecto tenga suficiente cuota de TPUS_PER_TPU_FAMILY, que especifica el número máximo de chips a los que puedes acceder en tuGoogle Cloud proyecto.
• Comprueba que tu proyecto tenga suficiente cuota de TPU para lo siguiente:
  • Cuota de máquinas virtuales de TPU
  • Cuota de direcciones IP
  • Cuota de Hyperdisk Balanced
• Permisos de proyecto de usuario
  • Si usas GKE con XPK, consulta los permisos de la consola de Cloud en la cuenta de usuario o de servicio para saber qué permisos necesitas para ejecutar XPK.

Crear variables de entorno

En Cloud Shell, crea las siguientes variables de entorno:

export PROJECT_ID=your-project-id
export TPU_NAME=your-tpu-name
export ZONE=us-central2-b
export ACCELERATOR_TYPE=v6e-4
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=your-service-account
export QUEUED_RESOURCE_ID=your-queued-resource-id

Descripciones de las variables de entorno

Variable	Descripción
PROJECT_ID	El ID de tu proyecto Google Cloud . Usa un proyecto que ya tengas o crea uno.
TPU_NAME	El nombre de la TPU.
ZONE	La zona en la que se creará la VM de TPU. Para obtener más información sobre las zonas admitidas, consulta Regiones y zonas de TPU.
ACCELERATOR_TYPE	El tipo de acelerador especifica la versión y el tamaño de la TPU de Cloud que quieres crear. Para obtener más información sobre los tipos de aceleradores compatibles con cada versión de TPU, consulta Versiones de TPU.
RUNTIME_VERSION	La versión de software de la TPU de Cloud.
SERVICE_ACCOUNT	La dirección de correo de tu cuenta de servicio. Para encontrarlo, ve a la página Cuentas de servicio de la consola Google Cloud . Por ejemplo: tpu-service-account@PROJECT_ID.iam.gserviceaccount.com
QUEUED_RESOURCE_ID	ID de texto asignado por el usuario de la solicitud de recurso en cola.

Aprovisionar una TPU v6e

    gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
        --node-id ${TPU_NAME} \
        --project ${PROJECT_ID} \
        --zone ${ZONE} \
        --accelerator-type ${ACCELERATOR_TYPE} \
        --runtime-version ${RUNTIME_VERSION} \
        --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 ver una lista completa de los estados de las solicitudes de recursos en cola, consulta la documentación sobre recursos en cola.

Conectarse a la TPU mediante SSH

  gcloud compute tpus tpu-vm ssh ${TPU_NAME}

Ejecutar la comparativa de Llama2-7B de PyTorch de JetStream

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

Cuando se haya completado la prueba de rendimiento de la inferencia, asegúrate de eliminar los recursos de TPU.

Limpieza

Elimina la TPU:

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