Diese Seite wurde von der Cloud Translation API übersetzt.

JetStream MaxText-Inferenz auf v6e-TPU-VMs

In dieser Anleitung erfahren Sie, wie Sie mit JetStream MaxText-Modelle auf TPU v6e bereitstellen. JetStream ist eine durchsatz- und speicheroptimierte Engine für die LLM-Inferenz (Large Language Model) auf XLA-Geräten (TPUs). In dieser Anleitung führen Sie den Inferenz-Benchmark für das Llama2-7B-Modell aus.

Hinweise

TPU v6e mit 4 Chips vorbereiten:

Melden Sie sich in Ihrem Google-Konto an. Wenn Sie noch kein Konto haben, melden Sie sich hier für ein neues Konto an.
Wählen Sie in der Google Cloud Console auf der Seite „Projektauswahl“ ein Google Cloud -Projekt aus oder erstellen Sie ein solches Projekt.
Aktivieren Sie die Abrechnung für Ihr Google Cloud -Projekt. Die Abrechnung ist für die gesamte Nutzung von Google Cloud erforderlich.
Installieren Sie die gcloud alpha-Komponenten.
Führen Sie den folgenden Befehl aus, um die neueste Version der gcloud-Komponenten zu installieren.
```
gcloud components update
```
Aktivieren Sie die TPU API mit dem folgenden gcloud-Befehl in Cloud Shell. Sie können sie auch über die Google Cloud Console aktivieren.
```
gcloud services enable tpu.googleapis.com
```

Erstellen Sie eine Dienstidentität für die TPU-VM.

gcloud alpha compute tpus tpu-vm service-identity create --zone=ZONE

Erstellen Sie ein TPU-Dienstkonto und gewähren Sie Zugriff auf Google Cloud- Dienste.

Dienstkonten ermöglichen dem Google Cloud- TPU-Dienst, auf andere Google Cloud--Dienste zuzugreifen. Ein nutzerverwaltetes Dienstkonto wird empfohlen. Folgen Sie diesen Anleitungen, um Rollen zu erstellen und zu gewähren. Folgende Rollen sind erforderlich:
- TPU-Administrator: Erforderlich, um eine TPU zu erstellen
- Storage Admin: Erforderlich für den Zugriff auf Cloud Storage
- Logs Writer: Erforderlich zum Schreiben von Protokollen mit der Logging API
- Monitoring Metric Writer: Erforderlich zum Schreiben von Messwerten in Cloud Monitoring
Authentifizieren Sie sich bei Google Cloud und konfigurieren Sie das Standardprojekt und die Standardzone für die Google Cloud CLI.
```
gcloud auth login
gcloud config set project PROJECT_ID
gcloud config set compute/zone ZONE
```

Sichere Kapazität

Wenden Sie sich an Ihr Cloud TPU-Vertriebs- oder Account-Team, um ein TPU-Kontingent anzufordern und Fragen zur Kapazität zu stellen.

Cloud TPU-Umgebung bereitstellen

Sie können v6e-TPUs mit GKE, mit GKE und XPK oder als Warteschlangenressourcen bereitstellen.

Vorbereitung

Prüfen Sie, ob für Ihr Projekt genügend TPUS_PER_TPU_FAMILY-Kontingent vorhanden ist. Dieses gibt die maximale Anzahl von Chips an, auf die Sie in IhremGoogle Cloud -Projekt zugreifen können.
Diese Anleitung wurde mit der folgenden Konfiguration getestet:
- Python 3.10 or later
- Nightly-Softwareversionen:
  - tägliche JAX 0.4.32.dev20240912
  - nightly LibTPU 0.1.dev20240912+nightly
- Stabile Softwareversionen:
  - JAX + JAX-Bibliothek von v0.4.35
Prüfen Sie, ob Ihr Projekt ein ausreichendes TPU-Kontingent für Folgendes hat:
- TPU-VM-Kontingent
- Kontingent für IP-Adressen
- Hyperdisk Balanced-Kontingent
Nutzerberechtigungen für Projekte
- Wenn Sie GKE mit XPK verwenden, finden Sie unter Cloud Console-Berechtigungen für das Nutzer- oder Dienstkonto Informationen zu den Berechtigungen, die zum Ausführen von XPK erforderlich sind.

Umgebungsvariablen erstellen

Erstellen Sie in Cloud Shell die folgenden Umgebungsvariablen:

export NODE_ID=TPU_NODE_ID # TPU name
export PROJECT_ID=PROJECT_ID
export ACCELERATOR_TYPE=v6e-4
export ZONE=us-central2-b
export RUNTIME_VERSION=v2-alpha-tpuv6e
export SERVICE_ACCOUNT=YOUR_SERVICE_ACCOUNT
export QUEUED_RESOURCE_ID=QUEUED_RESOURCE_ID
export VALID_DURATION=VALID_DURATION

# Additional environment variable needed for Multislice:
export NUM_SLICES=NUM_SLICES

# Use a custom network for better performance as well as to avoid having the
# default network becoming overloaded.
export NETWORK_NAME=${PROJECT_ID}-mtu9k
export NETWORK_FW_NAME=${NETWORK_NAME}-fw

Beschreibung der Befehls-Flags

Variable	Beschreibung
NODE_ID	Die vom Nutzer zugewiesene ID der TPU, die erstellt wird, wenn die anstehende Ressourcenanfrage zugewiesen wird.
PROJECT_ID	Google Cloud -Projektname. Verwenden Sie ein vorhandenes Projekt oder erstellen Sie ein neues.
ZONE	Welche Zonen unterstützt werden, erfahren Sie im Dokument TPU-Regionen und ‑Zonen.
ACCELERATOR_TYPE	Eine Liste aller unterstützten Beschleunigertypen finden Sie in der Dokumentation unter Beschleunigertypen.
RUNTIME_VERSION	`v2-alpha-tpuv6e`
SERVICE_ACCOUNT	Das ist die E-Mail-Adresse Ihres Dienstkontos. Sie finden sie in der Google Cloud Console unter „IAM“ -> „Dienstkonten“. Beispiel: tpu-service-account@<your_project_ID>.iam.gserviceaccount.com.com
NUM_SLICES	Die Anzahl der zu erstellenden Scheiben (nur für Mehrfachaufnahmen erforderlich)
QUEUED_RESOURCE_ID	Die vom Nutzer zugewiesene Text-ID der anstehenden Ressourcenanfrage.
VALID_DURATION	Die Dauer, für die die angeforderte Ressource gültig ist.
NETWORK_NAME	Der Name eines sekundären Netzwerks, das verwendet werden soll.
NETWORK_FW_NAME	Der Name einer sekundären Netzwerk-Firewall, die verwendet werden soll.

TPU v6e bereitstellen

gcloud alpha compute tpus queued-resources create QUEUED_RESOURCE_ID \
    --node-id TPU_NAME \
    --project PROJECT_ID \
    --zone ZONE \
    --accelerator-type v6e-4 \
    --runtime-version v2-alpha-tpuv6e \
    --service-account SERVICE_ACCOUNT

Verwenden Sie die Befehle list oder describe, um den Status der in der Warteschlange befindlichen Ressource abzufragen.

   gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}  \
      --project ${PROJECT_ID} --zone ${ZONE}

Eine vollständige Liste der Status von anstehenden Ressourcenanfragen finden Sie in der Dokumentation zu anstehenden Ressourcen.

Über SSH eine Verbindung zur TPU herstellen

   gcloud compute tpus tpu-vm ssh TPU_NAME

Sobald Sie eine Verbindung zur TPU hergestellt haben, können Sie den Inferenz-Benchmark ausführen.

Llama2-7B-Inferenz-Benchmark ausführen

Folgen Sie der Anleitung im GitHub-Repository, um JetStream und MaxText einzurichten, die Modell-Checkpunkte zu konvertieren und den Inferenz-Benchmark auszuführen.

Wenn der Inferenz-Benchmark abgeschlossen ist, bereinigen Sie die TPU-Ressourcen.

Bereinigen

Löschen Sie die TPU:

gcloud compute tpus queued-resources delete QUEUED_RESOURCE_ID \
    --project PROJECT_ID \
    --zone ZONE \
    --force \
    --async