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

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:

Sigue las instrucciones de la guía Configura el entorno de Cloud TPU para configurar un proyecto de Google Cloud , configurar Google Cloud CLI, habilitar la API de Cloud TPU y asegurarte de tener acceso para usar las Cloud TPU.
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, consulta Cuotas de Cloud TPU para obtener más información sobre las cuotas de Cloud TPU. Si tienes más preguntas sobre cómo proteger 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

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 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 marcas de comandos

Variable	Descripción
`PROJECT_ID`	Google Cloud nombre del proyecto. Usa un proyecto existente o crea uno nuevo.
`TPU_NAME`	El nombre de la TPU.
`ZONE`	Consulta el documento Regiones y zonas de TPU para conocer las zonas compatibles.
`ACCELERATOR_TYPE`	El tipo de acelerador especifica la versión y el tamaño de la Cloud TPU que deseas 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 Cloud TPU
`SERVICE_ACCOUNT`	La dirección de correo electrónico de tu cuenta de servicio . Para encontrarla, ve a la página Cuentas de servicio en la consola de Google Cloud. Por ejemplo: `tpu-service-account@PROJECT_ID.iam.gserviceaccount.com`
`QUEUED_RESOURCE_ID`	El ID de texto asignado por el usuario de la solicitud de recursos en cola.

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 ${ACCELERATOR_TYPE} \
        --runtime-version ${RUNTIME_VERSION} \
        --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