En esta guía se explica cómo servir un modelo de lenguaje extenso (LLM) mediante unidades de procesamiento tensorial (TPUs) en Google Kubernetes Engine (GKE) con JetStream a través de PyTorch. En esta guía, descargarás pesos de modelos en Cloud Storage y los desplegarás en un clúster Autopilot o Estándar de GKE mediante un contenedor que ejecute JetStream.
Si necesitas la escalabilidad, la resiliencia y la rentabilidad que ofrecen las funciones de Kubernetes al implementar tu modelo en JetStream, esta guía es un buen punto de partida.
Esta guía está dirigida a clientes de IA generativa que usan PyTorch, usuarios nuevos o actuales de GKE, ingenieros de aprendizaje automático, ingenieros de MLOps (DevOps) o administradores de plataformas que estén interesados en usar las funciones de orquestación de contenedores de Kubernetes para ofrecer LLMs.
Fondo
Si sirves un LLM mediante TPUs en GKE con JetStream, puedes crear una solución de servicio estable y lista para producción con todas las ventajas de Kubernetes gestionado, como la rentabilidad, la escalabilidad y la mayor disponibilidad. En esta sección se describen las tecnologías clave que se usan en este tutorial.
Acerca de las TPUs
Las TPUs son circuitos integrados para aplicaciones específicas (ASIC) desarrollados a medida por Google que se utilizan para acelerar los modelos de aprendizaje automático y de IA creados con frameworks como TensorFlow, PyTorch y JAX.
Antes de usar las TPUs en GKE, te recomendamos que completes el siguiente plan de formación:
- Consulta la arquitectura del sistema de las TPU de Cloud para obtener información sobre la disponibilidad de las versiones actuales de las TPU.
- Consulta información sobre las TPUs en GKE.
En este tutorial se explica cómo servir varios modelos de LLM. GKE implementa el modelo en nodos TPU v5e de un solo host con topologías de TPU configuradas en función de los requisitos del modelo para servir peticiones con baja latencia.
Acerca de JetStream
JetStream es un framework de servicio de inferencia de código abierto desarrollado por Google. JetStream permite realizar inferencias de alto rendimiento, alto volumen de procesamiento y memoria optimizada en TPUs y GPUs. JetStream ofrece optimizaciones de rendimiento avanzadas, como la creación de minilotes continua, la optimización de la caché de valores de clave y las técnicas de cuantización, para facilitar la implementación de LLMs. JetStream permite que el servicio de TPUs de PyTorch/XLA y JAX consiga un rendimiento óptimo.
Creación de lotes continua
El procesamiento por lotes continuo es una técnica que agrupa dinámicamente las solicitudes de inferencia entrantes en lotes, lo que reduce la latencia y aumenta el rendimiento.
Cuantización de la caché de KV
La cuantización de la caché de clave-valor consiste en comprimir la caché de clave-valor que se usa en los mecanismos de atención, lo que reduce los requisitos de memoria.
Cuantización de pesos Int8
La cuantización de pesos Int8 reduce la precisión de los pesos del modelo de coma flotante de 32 bits a enteros de 8 bits, lo que conlleva un cálculo más rápido y un menor uso de memoria.
Para obtener más información sobre estas optimizaciones, consulta los repositorios de proyectos JetStream PyTorch y JetStream MaxText.
Acerca de PyTorch
PyTorch es un framework de aprendizaje automático de código abierto desarrollado por Meta y que ahora forma parte de la Linux Foundation. PyTorch proporciona funciones de alto nivel, como el cálculo de tensores y las redes neuronales profundas.