Inferencias en Cloud TPU

La publicación hace referencia al proceso de implementación de un modelo de aprendizaje automático entrenado en un entorno de producción, en el que se puede usar para la inferencia. La inferencia es compatible con la TPU v5e y versiones posteriores. Los SLO de latencia son una prioridad para la entrega.

En este documento, se analiza la entrega de un modelo en una TPU de host único. Las porciones de TPU con 8 chips o menos tienen una VM o un host de TPU, y se denominan TPU de host único. Para obtener información sobre la inferencia con varios hosts, consulta Cómo realizar inferencias con varios hosts usando Pathways.

Comenzar

Necesitas una cuenta y un proyecto de Google Cloud para usar Cloud TPU. Para obtener más información, consulta Configura un entorno de Cloud TPU.

Asegúrate de tener suficiente cuota para la cantidad de núcleos de TPU que planeas usar para la inferencia. La TPU v5e usa cuotas separadas para el entrenamiento y la entrega. Las cuotas específicas de la entrega para la TPU v5e son las siguientes:

  • Recursos de v5e a pedido: TPUv5 lite pod cores for serving per project per zone
  • Recursos de v5e interrumpibles: Preemptible TPU v5 lite pod cores for serving per project per zone

Para otras versiones de TPU, las cargas de trabajo de entrenamiento y de servicio usan la misma cuota. Para obtener más información, consulta Cuotas de Cloud TPU.

Entrega LLMs con vLLM

vLLM es una biblioteca de código abierto diseñada para la inferencia y la entrega rápidas de modelos de lenguaje grandes (LLM). Cloud TPU se integra con vLLM a través del complemento tpu-inference, que admite modelos de JAX y PyTorch. Para obtener más información, consulta el repositorio de GitHub de tpu-inference.

Si deseas ver ejemplos del uso de vLLM para entregar un modelo en TPUs, consulta lo siguiente:

Generación de perfiles

Después de configurar la inferencia, puedes usar los generadores de perfiles para analizar el rendimiento y la utilización de la TPU. Para obtener más información sobre la generación de perfiles, consulta los siguientes recursos: