Inferenz auf Cloud TPU ausführen

Der Begriff der Bereitstellung bezieht sich auf die Bereitstellung eines trainierten Modells für maschinelles Lernen in einer Produktionsumgebung, in der es für die Inferenz verwendet werden kann. TPU v5e und neuere Versionen unterstützen die Inferenz. Die Einhaltung von Latenz-Service-Level-Objectives (SLOs) hat für die Bereitstellung von Arbeitslasten Priorität.

In diesem Dokument wird die Bereitstellung eines Modells auf einer TPU mit einem einzelnen Host behandelt. TPU-Slices mit maximal acht Chips haben eine TPU-VM oder einen Host und werden als TPUs mit einem Host bezeichnet. Informationen zur Inferenz auf mehreren Hosts mit PyTorch finden Sie unter vLLM-TorchTPU und SGLang-TorchTPU. Informationen zur Inferenz auf mehreren Hosts mit JAX finden Sie unter Inferenz auf mehreren Hosts mit Pathways durchführen.

TPU einrichten

Unabhängig davon, welchen Stack Sie auswählen, benötigen Sie ein Google Cloud -Projekt mit einer TPU mit einem einzelnen Host, bevor Sie ein Modell bereitstellen können. Die folgenden Schritte gelten für jede Kombination aus Bibliothek und Framework:

  1. Projekt erstellen Sie benötigen ein Konto und ein Projekt in Google Cloud , um Cloud TPU zu verwenden. Weitere Informationen finden Sie unter Google Cloud -Projekt für TPUs einrichten.

  2. Kontingent prüfen Prüfen Sie, ob Sie ein ausreichendes Kontingent für die Anzahl der TPU-Kerne haben, die Sie für die Inferenz verwenden möchten. TPU v5e verwendet separate Kontingente für Training und Bereitstellung. In der Cloud TPU API gelten die folgenden bereitstellungsspezifischen Kontingente für TPU v5e:

    • On-Demand-Ressourcen vom Typ v5e: TPU v5 lite pod cores for serving per project per zone
    • Ressourcen auf Abruf vom Typ v5e: Preemptible TPU v5 lite pod cores for serving per project per zone

    Bei anderen TPU-Versionen verwenden Trainings- und Bereitstellungsarbeitslasten dasselbe Kontingent. Weitere Informationen finden Sie unter Cloud TPU-Kontingente.

  3. Erstellen Sie eine TPU-VM. Wählen Sie eine TPU-Version und eine Zone mit Bereitstellungskapazität aus und erstellen Sie dann eine TPU-VM mit einem einzelnen Host. Weitere Informationen finden Sie unter TPU-Regionen und -Zonen und TPU-Ressourcen verwalten.

Nachdem die TPU-VM ausgeführt wird, folgen Sie der Einrichtungsanleitung für Ihre Bibliothek und Ihr Framework unter Inferenz-Stack auswählen.

Inferenz-Stack auswählen

Cloud TPU unterstützt die Inferenz von Large Language Models (LLMs) über zwei Open-Source-Bereitstellungsbibliotheken: vLLM und SGLang. Beide Bibliotheken werden auf Cloud TPU mit einem JAX- oder PyTorch-Backend ausgeführt. Sie können also die Bibliothek auswählen, die am besten zu Ihrer Arbeitslast passt, unabhängig davon, in welchem Framework Ihr Modell geschrieben ist.

Wählen Sie zuerst die Bibliothek und dann das Backend aus, das dem Framework Ihres Modells entspricht, wie in der folgenden Tabelle gezeigt:

Auslieferungsbibliothek JAX-Backend PyTorch-Backend
vLLM vllm-jax vllm-torchtpu
SGLang sglang-jax sglang-torchtpu

Inferenz mit vLLM

vLLM ist eine Open-Source-Bibliothek für schnelle Inferenz und Bereitstellung von LLMs, die auf kontinuierlicher Batchverarbeitung und PagedAttention basiert. Cloud TPU funktioniert mit vLLM über das vllm-jax-Plug-in für JAX-Modelle und das vllm-torchtpu-Plug-in für PyTorch-Modelle.

Weitere Beispiele für die Verwendung von vLLM zum Bereitstellen eines Modells auf TPUs finden Sie in den folgenden Ressourcen:

vLLM mit JAX

Die vLLM TPU-Dokumentation umfasst die Installation, eine Kurzanleitung, unterstützte Modelle und Abstimmungsoptionen. Weitere Informationen finden Sie unter JAX-Modelle mit vLLM auf TPU bereitstellen.

vLLM mit PyTorch

Das vllm-torchtpu-Plug-in fügt vLLM TorchTPU-Kernel- und Laufzeitunterstützung hinzu. Weitere Informationen finden Sie unter PyTorch-Modelle mit vLLM auf TPU bereitstellen.

Inferenz mit SGLang

SGLang ist ein Open-Source-Bereitstellungs-Framework für Large Language Models und multimodale Modelle, das eine leistungsstarke Laufzeit mit einem strukturierten Generierungs-Frontend kombiniert. Auf Cloud TPUs wird SGLang für JAX-Modelle über das SGLang-JAX-Backend (sglang-jax) und für PyTorch-Modelle über das SGLang-TorchTPU-Backend (sglang-torchtpu) ausgeführt.

SGLang mit PyTorch

Informationen zur Installation, zum Serverstart und zu unterstützten Konfigurationen finden Sie unter PyTorch-Modelle mit der SGLang-TorchTPU-Laufzeit auf TPU bereitstellen.

SGLang mit JAX

Informationen zur Installation, zum Serverstart und zu unterstützten Konfigurationen finden Sie unter JAX-Modelle mit der SGLang-JAX-Laufzeit auf TPU bereitstellen.

Leistung der Profilinferenz

Nachdem Sie die Inferenz eingerichtet haben, können Sie mit Profilern die Leistung und TPU-Auslastung analysieren. Weitere Informationen zur Profilerstellung finden Sie in den folgenden Ressourcen: