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:
Folgen Sie der Anleitung Cloud TPU-Umgebung einrichten, um sicherzustellen, dass Sie über den erforderlichen Zugriff für die Verwendung von Cloud TPUs verfügen.
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. Sie können ein Dienstkonto in der Google Cloud Console oder über den Befehl
gcloud
erstellen.Erstellen Sie ein Dienstkonto mit dem
gcloud
-Befehlszeilentool:gcloud iam service-accounts create your-service-account-name \ --description="your-sa-description" \ --display-name="your-sa-display-name" export SERVICE_ACCOUNT_NAME=your-service-account-name
So erstellen Sie ein Dienstkonto über die Google Cloud Console:
- Rufen Sie in der Google Cloud Console die Seite „Dienstkonten“ auf.
- Klicken Sie auf Dienstkonto erstellen.
- Geben Sie den Namen des Dienstkontos ein.
- Optional: Geben Sie eine Beschreibung für das Dienstkonto ein.
- Klicken Sie auf Erstellen und fahren Sie fort.
- Wählen Sie die Rollen aus, die Sie dem Dienstkonto zuweisen möchten.
- Klicken Sie auf Weiter.
- Optional: Geben Sie Nutzer oder Gruppen an, die das Dienstkonto verwalten können.
- Klicken Sie auf Fertig, um das Erstellen des Dienstkontos abzuschließen.
Nachdem Sie das Dienstkonto erstellt haben, können Sie ihm Rollen gewähren.
Folgende Rollen sind erforderlich:
- TPU-Administrator: Erforderlich zum Erstellen einer TPU
- Storage Admin: Erforderlich für den Zugriff auf Cloud Storage
- Log-Autor
- Monitoring-Messwert-Autor: Erforderlich zum Schreiben von Messwerten in Cloud Monitoring
Ihr Administrator muss Ihnen die Rolle
roles/resourcemanager.projectIamAdmin
zuweisen, damit Sie Nutzern IAM-Rollen zuweisen können. Auch ein Nutzer mit der Rolle „Projekt-IAM-Administrator“roles/resourcemanager.projectIamAdmin
kann diese Rolle gewähren.Verwenden Sie die folgenden
gcloud
-Befehle, um Dienstkontorollen hinzuzufügen:gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/tpu.admin gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/storage.admin gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/logging.logWriter gcloud projects add-iam-policy-binding ${PROJECT_ID} \ --member serviceAccount:${SERVICE_ACCOUNT_NAME}@${PROJECT_ID}.iam.gserviceaccount.com \ --role roles/monitoring.metricWriter
Sie können Rollen auch über die Google Cloud Console zuweisen.
Wählen Sie in der Google Cloud Console die folgenden Rollen aus:
- Wählen Sie Ihr Dienstkonto aus und klicken Sie auf Hauptkonto hinzufügen.
- Geben Sie im Feld Neue Hauptkonten die E-Mail-Adresse Ihres Dienstkontos ein.
- Suchen Sie im Drop-down-Menü Rolle auswählen nach einer Rolle (z. B. Storage-Administrator) und wählen Sie sie aus.
- Klicken Sie auf Speichern.
Authentifizieren Sie sich mit 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
Wenn Sie TPU-Kapazitäten sichern möchten, lesen Sie den Hilfeartikel Kontingente. Wenn Sie weitere Fragen zur Kapazitätssicherung haben, wenden Sie sich an Ihr Cloud TPU-Vertriebs- oder Account-Team.
Cloud TPU-Umgebung bereitstellen
Sie können TPU-VMs mit GKE, mit GKE und XPK oder als in der Warteschlange befindliche Ressourcen bereitstellen.
Vorbereitung
- Diese Anleitung wurde mit Python 3.10 oder höher getestet.
- Prüfen Sie, ob Ihr Projekt ein ausreichendes
TPUS_PER_TPU_FAMILY
-Kontingent hat. Dieses gibt die maximale Anzahl von Chips an, auf die Sie in IhremGoogle Cloud -Projekt zugreifen können. - 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 ist der 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 in der Warteschlange 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-Prüfpunkte 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