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:

Accede a tu Cuenta de Google. Si aún no lo hiciste, regístrate para obtener una nueva cuenta.
En la consola de Google Cloud, selecciona o crea un proyecto de Google Cloud en la página del selector de proyectos.
Habilita la facturación para tu proyecto de Google Cloud . La facturación es obligatoria para todo el uso de Google Cloud .
Instala los componentes de gcloud alpha.
Ejecuta el siguiente comando para instalar la versión más reciente de los componentes de gcloud.
```
gcloud components update
```
Habilita la API de TPU con el siguiente comando de gcloud en Cloud Shell. También puedes habilitarlo desde la consola de Google Cloud.
```
gcloud services enable tpu.googleapis.com
```

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 TPU de Google Cloud acceda a otros servicios de Google Cloud. Se recomienda una cuenta de servicio administrada por el usuario. Sigue estas guías para crear y otorgar roles. Se requieren los siguientes roles:
- Administrador de TPU: Necesario para crear una TPU
- Administrador de almacenamiento: Es necesario para acceder a Cloud Storage.
- Logs Writer: Es necesario para escribir registros con la API de Logging.
- Monitoring Metric Writer: Es necesario para escribir métricas en Cloud Monitoring.
Realiza la autenticación 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

Comunícate con tu equipo de ventas o de cuentas de Cloud TPU para solicitar la cuota de TPU y hacer preguntas sobre la capacidad.

Aprovisiona el entorno de Cloud TPU

Puedes aprovisionar TPUs v6e 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 .
Este instructivo se probó con la siguiente configuración:
- Python 3.10 or later
- Versiones de software nocturnas:
  - JAX nocturno 0.4.32.dev20240912
  - LibTPU nocturna 0.1.dev20240912+nightly
- Versiones de software estables:
  - JAX + JAX Lib de v0.4.35
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

Descripciones de las marcas de comandos

Variable	Descripción
NODE_ID	El ID asignado por el usuario de la TPU que se crea cuando se asigna la solicitud de recurso en fila.
ID DEL PROYECTO	El nombre del proyecto deGoogle Cloud Usa un proyecto existente o crea uno nuevo.
ZONA	Consulta el documento Regiones y zonas de TPU para conocer las zonas compatibles.
ACCELERATOR_TYPE	Consulta la documentación sobre Tipos de aceleradores para conocer los tipos de aceleradores compatibles.
RUNTIME_VERSION	`v2-alpha-tpuv6e`
SERVICE_ACCOUNT	Esta es la dirección de correo electrónico de tu cuenta de servicio que puedes encontrar en consola de Google Cloud -> IAM -> Cuentas de servicio Por ejemplo: tpu-service-account@<your_project_ID>.iam.gserviceaccount.com.com
NUM_SLICES	Es la cantidad de rebanadas que se deben crear (solo es necesario para Multislice).
QUEUED_RESOURCE_ID	El ID de texto asignado por el usuario de la solicitud de recursos en cola.
VALID_DURATION	Es la duración durante la cual la solicitud de recursos en cola es válida.
NETWORK_NAME	Es el nombre de una red secundaria que se usará.
NETWORK_FW_NAME	Es el nombre de un firewall de red secundario que se usará.

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