Diffuser des LLM à l'aide de TPU multi-hôtes sur GKE avec JetStream et Pathways

Ce guide vous explique comment diffuser des grands modèles de langage (LLM) de pointe tels que Llama 3.1 405B sur Google Kubernetes Engine (GKE) à l'aide d'unités de traitement tensoriel (TPU) sur plusieurs nœuds.

Ce guide explique comment utiliser des technologies Open Source portables (Kubernetes, JetStream, Pathways on Cloud et l'API LeaderWorkerSet (LWS)) pour déployer et diffuser des charges de travail d'IA/de ML sur GKE, en tirant parti du contrôle précis, de l'évolutivité, de la résilience, de la portabilité et de la rentabilité de GKE.

Arrière-plan

La taille des grands modèles de langage a augmenté et ils ne tiennent plus sur une seule tranche de TPU hôte. Pour l'inférence ML, vous pouvez utiliser Pathways on Cloud afin d'exécuter une inférence multi-hôtes à grande échelle sur GKE sur plusieurs nœuds TPU interconnectés. Dans ce guide, vous allez apprendre à provisionner un cluster GKE avec les tranches TPU multi-hôtes, à utiliser les binaires Pathways on Cloud, à lancer le serveur JetStream avec le framework MaxText et à effectuer des requêtes d'inférence multi-hôtes.

En diffusant un LLM à l'aide de TPU sur GKE avec JetStream, MaxText et Pathways, vous pouvez créer une solution de diffusion robuste et prête pour la production avec tous les avantages de la plate-forme Kubernetes gérée, y compris en termes de rentabilité, évolutivité et haute disponibilité. Cette section décrit les principales technologies utilisées dans ce tutoriel.

À propos des TPU

Les TPU sont des circuits intégrés propres aux applications (ASIC) développés spécifiquement par Google et permettant d'accélérer le machine learning et les modèles d'IA créés à l'aide de frameworks tels que TensorFlow, PyTorch et JAX.

Avant d'utiliser des TPU dans GKE, nous vous recommandons de suivre le parcours de formation suivant :

  1. Découvrez la disponibilité actuelle des versions de TPU avec l'architecture système de Cloud TPU.
  2. Apprenez-en plus sur les TPU dans GKE.

Ce tutoriel explique comment diffuser le modèle Llama 3.1-405B. GKE déploie le modèle sur des nœuds TPU v6e multi-hôtes avec des topologies TPU configurées en fonction des exigences du modèle pour diffuser des requêtes avec une faible latence.

Parcours sur Cloud

Pathways est une couche d'orchestration à grande échelle pour les accélérateurs. Pathways est explicitement conçu pour permettre l'exploration de nouveaux systèmes et de nouvelles idées de recherche en ML, tout en conservant des performances de pointe pour les modèles actuels. Pathways permet à un seul processus client JAX de coordonner le calcul sur une ou plusieurs grandes tranches de TPU, ce qui simplifie les calculs de ML qui s'étendent sur des centaines ou des milliers de puces TPU.

JetStream

JetStream est un framework de diffusion d'inférences Open Source développé par Google. JetStream permet des inférences hautes performances, à haut débit et à mémoire optimisée sur les TPU et les GPU. JetStream fournit des optimisations de performances avancées, y compris des techniques de traitement par lot continu, d'optimisation du cache KV et de quantification, pour faciliter le déploiement de LLM. JetStream permet d'optimiser les performances du service TPU PyTorch/XLA et JAX.

MaxText

MaxText est une implémentation LLM JAX performante, évolutive et adaptable, basée sur des bibliothèques JAX Open Source telles que Flax, Orbax et Optax. L'implémentation LLM uniquement décodeur de MaxText est écrite en Python. Elle exploite fortement le compilateur XLA pour atteindre de hautes performances sans avoir à créer de noyau personnalisé.

Pour en savoir plus sur les derniers modèles et tailles de paramètres compatibles avec MaxText, consultez le dépôt du projet MaxText.

Llama 3.1 405B

Llama 3.1 405B est un grand modèle de langage de Meta conçu pour diverses tâches de traitement du langage naturel, y compris la génération de texte, la traduction et les systèmes de questions-réponses. GKE fournit l'infrastructure nécessaire pour répondre aux besoins d'entraînement et de diffusion distribués des modèles de cette envergure.

Pour en savoir plus, consultez la documentation Llama.

Architecture

Cette section décrit l'architecture GKE utilisée dans ce tutoriel. L'architecture comprend un cluster GKE Standard qui provisionne des TPU et héberge des composants JetStream et Pathways pour déployer et diffuser le modèle.

Le schéma suivant montre les composants de cette architecture:

Architecture d'un cluster GKE avec un pool de nœuds TPU multi-hôte contenant les composants JetStream et Pathways.

Cette architecture comprend les composants suivants :

  • Un cluster régional GKE Standard.
  • Un pool de nœuds de tranche de TPU multi-hôte qui héberge le déploiement JetStream et les composants Pathways.
  • Pathways resource manager gère les ressources d'accélérateur et coordonne l'allocation des accélérateurs pour les jobs utilisateur.
  • Pathways client se coordonne avec Pathways resource manager pour déterminer où les programmes compilés sont placés pour l'exécution.
  • Le Pathways worker s'exécute et effectue des calculs sur les machines d'accélérateur, puis renvoie les données à votre charge de travail via le serveur proxy IFRT.
  • IFRT proxy client implémente l'API OSS Interim Framework Runtime (IFRT) et sert de pont de communication entre votre charge de travail et les composants Pathways.
  • Le IFRT proxy server reçoit les requêtes du IFRT proxy client et les transmet au Pathways client, en répartissant le travail.
  • Le conteneur JetStream-Pathways fournit un serveur d'inférence basé sur JAX qui reçoit les requêtes d'inférence et délègue ses processus d'exécution à Pathways workers.
  • Le composant de service répartit le trafic entrant sur toutes les instances répliquées JetStream HTTP.
  • JetStream HTTP est un serveur HTTP qui accepte les requêtes en tant qu'encapsuleur au format requis par JetStream et les envoie au client GRPC de JetStream.