Hex-LLM: Modelo de lenguaje grande de alta eficiencia que se ejecuta en TPUs en el Jardín de Modelos de Vertex AI

El jardín de modelos de Vertex AI de Google Cloud tiene como objetivo optimizar los flujos de trabajo de aprendizaje automático con más de 150 modelos, incluyendo modelos de primera y tercera partes, y de código abierto. El año pasado, introdujeron la pila de servicio vLLM en GPUs, y ahora presentan Hex-LLM, optimizado para Cloud TPUs utilizando XLA. Hex-LLM mejora la eficiencia y la rentabilidad al servir grandes modelos de lenguaje (LLMs) al integrar tecnologías de vanguardia como la agrupación continua y la atención paginada. Soporta una variedad de modelos de LLM densos y dispersos populares y ofrece un alto rendimiento y baja latencia. Las optimizaciones clave incluyen un algoritmo de agrupación continua basado en tokens, un kernel de PagedAttention reescrito y estrategias de paralelismo de datos y tensor flexible. Las pruebas de rendimiento con el conjunto de datos ShareGPT mostraron impresionantes métricas de rendimiento, con modelos como Gemma 7B y Llama 2 70B entregando resultados competitivos en chips TPU v5e. Los usuarios pueden implementar fácilmente Hex-LLM a través del jardín de modelos de Vertex AI, la implementación de un clic o las libretas de Colab Enterprise. Esta flexibilidad permite la personalización para manejar necesidades de tráfico variables, lo que hace que Hex-LLM sea una herramienta poderosa para el servicio eficiente de LLM en hardware TPU de Google.