In diesem Leitfaden erfahren Sie, wie Sie ein Large Language Model (LLM) mit Tensor Processing Units (TPUs) in Google Kubernetes Engine (GKE) mit JetStream über PyTorch bereitstellen. In dieser Anleitung laden Sie Modellgewichte in Cloud Storage herunter und stellen sie auf einem GKE-Autopilot oder -Standard mit einem Container, der JetStream ausführt, bereit.
Wenn Sie beim Bereitstellen Ihres Modells in JetStream die Skalierbarkeit, Robustheit und Kosteneffizienz der Kubernetes-Features benötigen, ist dieser Leitfaden ein guter Ausgangspunkt.
Dieser Leitfaden richtet sich an Kunden von generativer KI, die PyTorch verwenden, neue oder bestehende Nutzer von GKE, ML-Entwickler, MLOps-Entwickler (DevOps) oder Plattformadministratoren, die daran interessiert sind, Funktionen zur Kubernetes-Containerorchestrierung für die Bereitstellung von LLMs zu nutzen.
Hintergrund
Wenn Sie ein LLM mithilfe von TPUs in GKE mit JetStream bereitstellen, können Sie eine robuste, produktionsreife Bereitstellungslösung mit allen Vorteilen von verwaltetem Kubernetes erstellen, einschließlich Kosteneffizienz, Skalierbarkeit und höhere 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 folgenden erstellt wurden:TensorFlow, PyTorch und JAX.
Bevor Sie TPUs in GKE verwenden, sollten Sie den folgenden Lernpfad durcharbeiten:
- Lernen Sie mehr über die aktuelle Verfügbarkeit von TPU-Versionen unter Cloud TPU-Systemarchitektur.
- TPUs in GKE
In dieser Anleitung wird das Bereitstellen verschiedener LLM-Modelle beschrieben. GKE stellt das Modell auf TPUv5e-Knoten mit einem einzelnen Host bereit. Dabei werden TPU-Topologien basierend auf den Modellanforderungen für die Bereitstellung von Eingabeaufforderungen mit niedriger Latenz konfiguriert.
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 eine optimale Leistung zu erzielen.
Kontinuierliche Batchverarbeitung
Kontinuierliches Batching ist eine Technik, bei der eingehende Inferenzanfragen dynamisch in Batches gruppiert werden, um die Latenz zu verringern und den Durchsatz zu erhöhen.
KV-Cache-Quantisierung
Bei der KV-Cache-Quantisierung wird der im Attention-Mechanismus verwendete Schlüssel-Wert-Cache komprimiert, wodurch der Arbeitsspeicherbedarf sinkt.
Int8-Gewichtsquantisierung
Durch die Int8-Gewichtsquantisierung wird die Genauigkeit der Modellgewichte von 32-Bit-Gleitkommazahlen auf 8-Bit-Ganzzahlen reduziert, was zu schnelleren Berechnungen und einer geringeren Speichernutzung führt.
Weitere Informationen zu diesen Optimierungen finden Sie in den Projekt-Repositories JetStream PyTorch und JetStream MaxText.
PyTorch
PyTorch ist ein Open-Source-Framework für maschinelles Lernen, das von Meta entwickelt wurde und jetzt Teil der Linux Foundation ist. PyTorch bietet High-Level-Funktionen wie Tensor-Berechnungen und neuronale Deep-Learning-Netzwerke.