LLMs mithilfe von Multi-Host-TPUs in GKE mit JetStream und Pathways bereitstellen

In diesem Leitfaden erfahren Sie, wie Sie moderne Large Language Models (LLMs) wie Llama 3.1 405B in Google Kubernetes Engine (GKE) mit Tensor Processing Units (TPUs) auf mehreren Knoten bereitstellen.

In diesem Leitfaden wird gezeigt, wie Sie mit portablen Open-Source-Technologien wie Kubernetes, JetStream, Pathways on Cloud und der LeaderWorkerSet-API (LWS) KI/ML-Arbeitslasten in GKE bereitstellen und bereitstellen können. Dabei werden die detaillierte Steuerung, Skalierbarkeit, Robustheit, Portabilität und Kosteneffizienz von GKE genutzt.

Wichtig :Wenn Sie Zugriff auf Pathways erhalten möchten, wenden Sie sich an Ihren Google Cloud-Kundenbetreuer.

Hintergrund

Large Language Models sind immer größer geworden und passen nicht mehr auf ein einzelnes Host-TPU-Slice. Für die ML-Inferenz können Sie Pathways on Cloud verwenden, um die Inferenz mit mehreren Hosts in großem Maßstab in GKE auf mehreren miteinander verbundenen TPU-Knoten auszuführen. In dieser Anleitung erfahren Sie, wie Sie einen GKE-Cluster mit den Multi-Host-TPU-Slices bereitstellen, die Binärdateien von Pathways on Cloud verwenden, den JetStream-Server mit dem MaxText-Framework starten und Multi-Host-Inferenzanfragen stellen.

Wenn Sie ein LLM mit TPUs in GKE mit JetStream, MaxText und Pathways bereitstellen, können Sie eine robuste, produktionsreife Bereitstellungslösung mit allen Vorteilen von verwaltetem Kubernetes erstellen, einschließlich Kosteneffizienz, Skalierbarkeit und höherer Verfügbarkeit. In diesem Abschnitt werden die in dieser Anleitung verwendeten Schlüsseltechnologien beschrieben.

TPUs

TPUs sind von Google speziell entwickelte anwendungsspezifische integrierte Schaltungen (Application-Specific Integrated Circuits, ASICs), die verwendet werden, um das maschinelle Lernen und die KI-Modelle zu beschleunigen, die mit Frameworks wie TensorFlow, PyTorch und JAX erstellt wurden.

Bevor Sie TPUs in GKE verwenden, sollten Sie den folgenden Lernpfad durcharbeiten:

1. Lernen Sie mehr über die aktuelle Verfügbarkeit von TPU-Versionen unter Cloud TPU-Systemarchitektur.
2. TPUs in GKE

In dieser Anleitung wird das Bereitstellen des Modells Llama 3.1-405B beschrieben. GKE stellt das Modell auf TPU v6e-Knoten mit mehreren Hosts bereit. Dabei werden TPU-Topologien basierend auf den Modellanforderungen für die Bereitstellung von Prompts mit niedriger Latenz konfiguriert.

Pathways on Cloud

Pathways ist eine Orchestrierungsebene für Beschleuniger im großen Maßstab. Pathways wurde speziell entwickelt, um die Erforschung neuer Systeme und ML-Forschungsideen zu ermöglichen und gleichzeitig die Leistung aktueller Modelle auf dem neuesten Stand zu halten. Mit Pathways kann ein einzelner JAX-Clientprozess Berechnungen über einen oder mehrere große TPU-Slices hinweg koordinieren. So werden ML-Berechnungen optimiert, die sich über Hunderte oder Tausende von TPU-Chips erstrecken.

JetStream

JetStream ist ein von Google entwickeltes Open-Source-Framework zur Bereitstellung von Inferenzen. JetStream ermöglicht leistungsstarke, durchsatzintensive und speicheroptimierte Inferenz auf TPUs und GPUs. JetStream bietet erweiterte Leistungsoptimierungen, einschließlich Techniken zur kontinuierlichen Batch-, KV-Cache- und Quantisierung, um die LLM-Bereitstellung zu erleichtern. JetStream ermöglicht die PyTorch/XLA- und JAX-TPU-Bereitstellung, um die Leistung zu optimieren.

MaxText

MaxText ist eine leistungsstarke, skalierbare und anpassbare JAX LLM-Implementierung, die auf Open-Source-JAX-Bibliotheken auf folgende Weise basiert: Flax, Orbax undOptax. Die nur Decoder-LLM-Implementierung von MaxText ist in Python geschrieben. Es nutzt den XLA-Compiler stark aus, um eine hohe Leistung zu erzielen, ohne benutzerdefinierte Kernel erstellen zu müssen.

Weitere Informationen zu den neuesten Modellen und Parametergrößen, die von MaxText unterstützt werden, finden Sie im Projekt-Repository MaxText.

Llama 3.1 405B

Llama 3.1 405B ist ein Large Language Model von Meta, das für eine Reihe von Natural Language Processing-Aufgaben entwickelt wurde, darunter Textgenerierung, Übersetzung und Question Answering. GKE bietet die Infrastruktur, die für das verteilte Training und die Bereitstellung von Modellen dieser Größenordnung erforderlich ist.

Weitere Informationen finden Sie in der Llama-Dokumentation.

Architektur

In diesem Abschnitt wird die in dieser Anleitung verwendete GKE-Architektur beschrieben. Die Architektur umfasst einen GKE-Standardcluster, der TPUs bereitstellt und JetStream- und Pathways-Komponenten zum Bereitstellen und Bereitstellen des Modells hostet.

Das folgende Diagramm zeigt die Komponenten dieser Architektur:

Diese Architektur umfasst die folgenden Komponenten:

• Einen regionalen GKE Standard.
• Ein TPU-Slice-Knotenpool mit mehreren Hosts, der die JetStream-Bereitstellung und Pathways-Komponenten hostet.
• Der Pathways resource manager verwaltet Beschleunigerressourcen und koordiniert die Zuweisung von Beschleunigern für Nutzerjobs.
• Die Pathways client-Koordinaten werden mit Pathways resource manager abgestimmt, um festzulegen, wo die kompilierten Programme zur Ausführung platziert werden.
• Die Pathways worker wird auf Beschleuniger-Maschinen ausgeführt und führt Berechnungen durch. Anschließend werden Daten über den IFRT-Proxyserver an Ihre Arbeitslast zurückgesendet.
• Das IFRT proxy client implementiert die OSS-API Interim Framework Runtime (IFRT) und fungiert als Kommunikationsbrücke zwischen Ihrer Arbeitslast und den Pathways-Komponenten.
• Der IFRT proxy server empfängt Anfragen vom IFRT proxy client und leitet sie an den Pathways client weiter, um die Arbeit zu verteilen.
• Der JetStream-Pathways-Container stellt einen JAX-basierten Inferenzserver bereit, der Inferenzanfragen empfängt und seine Ausführungsprozesse an den Pathways workers-Container delegiert.
• Die Dienstkomponente verteilt eingehenden Traffic auf alle JetStream HTTP-Replikate.
• JetStream HTTP ist ein HTTP-Server, der Anfragen als Wrapper für das erforderliche Format von JetStream akzeptiert und an den GRPC-Client von JetStream sendet.