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

Inferencia de JetStream MaxText en VMs de TPU v6e

En este instructivo, se muestra cómo usar JetStream para entregar modelos de MaxText 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 la guía Configura el entorno de Cloud TPU para asegurarte de tener el acceso adecuado para usar Cloud TPU.

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

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	Google Cloud nombre del proyecto. 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 de Tipos de aceleradores para conocer todos 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 la 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

Una vez que te conectes a la TPU, podrás ejecutar la comparativa de inferencia.

Ejecuta la comparativa de inferencia Llama2-7B

Para configurar JetStream y MaxText, 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