Hex-LLM : Modèle de langage grand et efficace servant sur des TPU dans le Jardin de Modèles Vertex AI

Le jardin de modèles Vertex AI de Google Cloud vise à optimiser les flux de travail d'apprentissage automatique avec plus de 150 modèles, notamment des modèles de première partie, open-source et tiers. L'année dernière, ils ont introduit la pile de service vLLM sur les GPU, et maintenant ils dévoilent Hex-LLM, optimisé pour les Cloud TPUs en utilisant XLA. Hex-LLM améliore l'efficacité et la rentabilité dans la mise en service de grands modèles de langage (LLM) en intégrant des technologies de pointe telles que le batchage continu et l'attention paginée. Il prend en charge une gamme de modèles de langage dense et sparse populaires et offre un débit élevé et une faible latence. Les optimisations clés comprennent un algorithme de batchage continu basé sur les jetons, un noyau PagedAttention réécrit et des stratégies de parallélisme de données et de tenseurs flexibles. Les benchmarks avec le jeu de données ShareGPT ont montré des métriques de performance impressionnantes, avec des modèles comme Gemma 7B et Llama 2 70B offrant des résultats compétitifs sur les puces TPU v5e. Les utilisateurs peuvent déployer facilement Hex-LLM via le playground du jardin de modèles Vertex AI, le déploiement en un clic ou les notebooks Colab Enterprise. Cette flexibilité permet une personnalisation pour gérer les besoins de trafic variables, ce qui fait de Hex-LLM un outil puissant pour la mise en service efficace des LLM sur le matériel TPU de Google.