En esta guía, se muestra cómo entregar modelos de lenguaje grandes (LLMs) de vanguardia, como Llama 3.1 405B, en Google Kubernetes Engine (GKE) con unidades de procesamiento tensorial (TPUs) en varios nodos.
En esta guía, se muestra cómo usar tecnologías portátiles de código abierto (Kubernetes, JetStream, Pathways on Cloud y la API de LeaderWorkerSet [LWS]) para implementar y entregar cargas de trabajo de IA/AA en GKE, aprovechando el control detallado, la escalabilidad, la resiliencia, la portabilidad y la rentabilidad de GKE.
Fondo
Los modelos de lenguaje grandes crecieron en tamaño y ya no caben en una sola porción de TPU de host. Para la inferencia de AA, puedes usar Pathways en Cloud para ejecutar inferencias a gran escala con varios hosts en GKE en varios nodos de TPU interconectados. En esta guía, aprenderás a aprovisionar un clúster de GKE con las porciones de TPU de varios hosts y a usar los archivos binarios de Pathways on Cloud, iniciar el servidor de JetStream con el framework de MaxText y realizar solicitudes de inferencia de varios hosts.
Con la entrega de un LLM con TPU en GKE con JetStream, MaxText y Pathways, puedes compilar una solución de entrega sólida y lista para la producción con todos los beneficios de Kubernetes administrado, incluida la rentabilidad, escalabilidad y disponibilidad mayor. En esta sección, se describen las tecnologías clave que se usan en este instructivo.
Acerca de las TPU
Las TPU son circuitos integrados personalizados específicos de aplicaciones (ASIC) de Google que se usan para acelerar el aprendizaje automático y los modelos de IA compilados con frameworks como TensorFlow, PyTorch y JAX.
Antes de usar las TPU en GKE, te recomendamos que completes la siguiente ruta de aprendizaje:
- Obtén información sobre la disponibilidad actual de la versión de TPU con la arquitectura del sistema de Cloud TPU.
- Obtén información sobre las TPU en GKE.
En este instructivo, se aborda la entrega del modelo Llama 3.1-405B. GKE implementa el modelo en los nodos TPU v6e de varios hosts con topologías de TPU configuradas según los requisitos del modelo para entregar mensajes con baja latencia.
Rutas de aprendizaje en Cloud
Pathways es una capa de organización a gran escala para aceleradores. Pathways está diseñado explícitamente para permitir la exploración de nuevos sistemas e ideas de investigación de AA, al tiempo que conserva el rendimiento de vanguardia para los modelos actuales. Pathways permite que un solo proceso cliente de JAX coordine el procesamiento en una o más porciones de TPU grandes, lo que optimiza los cálculos de AA que abarcan cientos o miles de chips de TPU.
JetStream
JetStream es un framework de entrega de inferencia de código abierto que desarrolla Google. JetStream permite la inferencia de alto rendimiento, alta capacidad de procesamiento y con optimización de memoria en TPU y GPU. JetStream proporciona optimizaciones de rendimiento avanzadas, incluidas técnicas de procesamiento por lotes, optimizaciones de la caché de KV y de cuantización continuas, para facilitar la implementación de LLM. JetStream permite que PyTorch/XLA y JAX TPU entreguen un rendimiento óptimo.
MaxText
MaxText es una implementación de LLM JAX escalable y adaptable, compilada en bibliotecas JAX de código abierto como Flax, Orbax y Optax. La implementación de LLM solo de decodificador de MaxText se escribe en Python. Aprovecha en gran medida el compilador XLA para lograr un alto rendimiento sin necesidad de compilar kernels personalizados.
Para obtener más información sobre los modelos y tamaños de parámetros más recientes que admite MaxText, consulta el repositorio del proyecto de MaxText.
Llama 3.1 405b
Llama 3.1 405B es un modelo de lenguaje grande de Meta diseñado para una variedad de tareas de procesamiento de lenguaje natural, como la generación de texto, la traducción y la respuesta de preguntas. GKE ofrece la infraestructura necesaria para admitir las necesidades de entrenamiento y entrega distribuidos de modelos de esta escala.
Para obtener más información, consulta la documentación de Llama.
Arquitectura
En esta sección, se describe la arquitectura de GKE que se usa en este instructivo. La arquitectura incluye un clúster de GKE Standard que aprovisiona TPU y aloja componentes de JetStream y Pathways para implementar y entregar el modelo.
En el siguiente diagrama, se muestran los componentes de esta arquitectura:
Esta arquitectura incluye los siguientes componentes:
- Un clúster regional de GKE Standard.
- Un grupo de nodos de porción de TPU de varios hosts que aloja los componentes de implementación y de Pathways de JetStream.
Pathways resource manageradministra los recursos del acelerador y coordina la asignación de aceleradores para los trabajos del usuario.- El objeto
Pathways clientse coordina con el objetoPathways resource managerpara determinar dónde se colocan los programas compilados para su ejecución. - El
Pathways workerse ejecuta y realiza cálculos en máquinas aceleradoras, y envía datos de vuelta a tu carga de trabajo a través del servidor proxy de IFRT. IFRT proxy clientimplementa la API de Interim Framework Runtime (IFRT) de OSS y actúa como puente de comunicación entre tu carga de trabajo y los componentes de Pathways.- El
IFRT proxy serverrecibe solicitudes delIFRT proxy clienty las reenvía alPathways client, distribuyendo el trabajo. - El contenedor
JetStream-Pathwaysproporciona un servidor de inferencia basado en JAX que recibe solicitudes de inferencia y delega sus procesos de ejecución enPathways workers. - El componente de Service distribuye el tráfico entrante a todas las réplicas de
JetStream HTTP. JetStream HTTPes un servidor HTTP que acepta solicitudes como wrapper del formato requerido de JetStream y las envía al cliente GRPC de JetStream.