Inferencia de TPU de Cloud

El servicio se refiere al proceso de desplegar un modelo de aprendizaje automático entrenado en un entorno de producción, donde se puede usar para la inferencia. La inferencia se admite en la versión 5e de TPU y en versiones posteriores. Los SLOs de latencia son una prioridad para el servicio.

En este documento se explica cómo publicar un modelo en una TPU de un solo host. Los sectores de TPU con 8 chips o menos tienen una VM o un host de TPU y se denominan TPUs de un solo host.

Empezar

Para usar Cloud TPU, necesitas una Google Cloud cuenta y un proyecto. Para obtener más información, consulta el artículo Configurar un entorno de TPU de Cloud.

Debes solicitar la siguiente cuota para publicar contenido en TPUs:

  • Recursos de v5e bajo demanda: TPUv5 lite pod cores for serving per project per zone
  • Recursos de v5e preemptivos: Preemptible TPU v5 lite pod cores for serving per project per zone
  • Recursos de la versión 6e bajo demanda: TPUv6 cores per project per zone
  • Recursos v6e de carácter temporal: Preemptible TPUv6 cores per project per zone

Para obtener más información sobre las cuotas de TPU, consulta Cuotas de TPU.

Servir LLMs con JetStream

JetStream es un motor optimizado de rendimiento y memoria para la inferencia de modelos de lenguaje extensos (LLMs) en dispositivos XLA (TPUs). Puedes usar JetStream con modelos de JAX y PyTorch/XLA. Para ver un ejemplo de cómo usar JetStream para servir un LLM de JAX, consulta Inferencia de MaxText de JetStream en TPU v6e.

Servir modelos LLM con vLLM

vLLM es una biblioteca de código abierto diseñada para la inferencia y el servicio rápidos de modelos de lenguaje extensos (LLMs). Puedes usar vLLM con PyTorch/XLA. Para ver un ejemplo de cómo usar vLLM para servir un LLM de PyTorch, consulta Sirve un LLM mediante TPU Trillium en GKE con vLLM.

Elaboración de perfiles

Una vez que hayas configurado la inferencia, puedes usar los profilers para analizar el rendimiento y el uso de las TPU. Para obtener más información sobre la creación de perfiles, consulta los siguientes artículos: