Hex-LLM: Modelo de linguagem grande de alta eficiência servido em TPUs no Jardim de Modelos Vertex AI

O Jardim de Modelos de IA Vertex da Google Cloud visa otimizar fluxos de trabalho de aprendizado de máquina com mais de 150 modelos, incluindo modelos de primeira e terceira partes e de código aberto. No ano passado, eles introduziram a pilha de serviço vLLM em GPUs, e agora revelam o Hex-LLM, otimizado para Cloud TPUs usando XLA. O Hex-LLM melhora a eficiência e a eficácia em termos de custo no serviço de grandes modelos de linguagem (LLMs) ao incorporar tecnologias de ponta como batching contínuo e atenção paginada. Ele suporta uma variedade de modelos de LLM densos e esparsos populares e oferece alto throughput e baixa latência. As principais otimizações incluem um algoritmo de batching contínuo baseado em tokens, um kernel de PagedAttention reescrito e estratégias de paralelismo de dados e tensor flexíveis. A benchmarking com o conjunto de dados ShareGPT mostrou impressionantes métricas de desempenho, com modelos como Gemma 7B e Llama 2 70B entregando resultados competitivos em chips TPU v5e. Os usuários podem implantar facilmente o Hex-LLM via playground do Jardim de Modelos de IA Vertex, implantação em um clique ou Notebooks Enterprise Colab. Essa flexibilidade permite personalizar para lidar com necessidades de tráfego variáveis, tornando o Hex-LLM uma ferramenta poderosa para o serviço eficiente de LLMs no hardware TPU da Google.