Dieses Dokument enthält eine kurze Einführung in die Arbeit mit Pax auf einer TPU mit einem einzelnen Host (v2-8, v3-8, v4-8).
Pax ist ein Framework zum Konfigurieren und Ausführen von Tests für maschinelles Lernen zusätzlich zu JAX. Pax konzentriert sich auf die Vereinfachung von ML in großem Maßstab durch die gemeinsame Nutzung von Infrastrukturkomponenten mit vorhandenen ML-Frameworks und die Nutzung der Praxis-Modellierungsbibliothek für Modularität.
Lernziele
- TPU-Ressourcen für das Training einrichten
- Pax auf einer TPU mit einem einzelnen Host installieren
- Transformer-basiertes SPMD-Modell mit Pax trainieren
Hinweise
Führen Sie die folgenden Befehle aus, um gcloud für die Verwendung Ihres Cloud TPU-Projekts zu konfigurieren und Komponenten zu installieren, die zum Trainieren eines Modells benötigt werden, das Pax auf einer TPU mit einem einzelnen Host ausführt.
Google Cloud CLI installieren
Die Google Cloud CLI enthält Tools und Bibliotheken für die Interaktion mit Produkten und Diensten der Google Cloud CLI. Wenn Sie es noch nicht installiert haben, installieren Sie es jetzt. Folgen Sie dazu der Anleitung unter Google Cloud CLI installieren.
gcloud-Befehl konfigurieren
Führen Sie den gcloud auth list aus, um die derzeit verfügbaren Konten zu sehen.
$ gcloud config set account account
$ gcloud config set project project-id
Cloud TPU API aktivieren
Aktivieren Sie die Cloud TPU API mit dem folgenden gcloud-Befehl in Cloud Shell.
Sie können es auch über die Google Cloud Console aktivieren.
$ gcloud services enable tpu.googleapis.com
Führen Sie den folgenden Befehl aus, um eine Dienstidentität (ein Dienstkonto) zu erstellen.
$ gcloud beta services identity create --service tpu.googleapis.com
TPU-VM erstellen
Mit Cloud TPU-VMs werden Modell und Code direkt auf der TPU-VM ausgeführt. Stellen Sie eine SSH-Verbindung zur TPU-VM her. Sie können beliebigen Code direkt auf der TPU-VM ausführen, Pakete installieren, Logs ansehen und Code debuggen.
Erstellen Sie die TPU-VM. Führen Sie dazu den folgenden Befehl in Cloud Shell oder auf dem Computerterminal aus, auf dem die Google Cloud CLI installiert ist.
Legen Sie die zone gemäß der Verfügbarkeit in Ihrem Vertrag fest und verweisen Sie bei Bedarf auf TPU-Regionen und -Zonen.
Legen Sie die Variable accelerator-type auf v2-8, v3-8 oder v4-8 fest.
Legen Sie die Variable version für v2- und v3-TPU-Versionen auf tpu-vm-base oder für v4-TPUs auf tpu-vm-v4-base fest.
$ gcloud compute tpus tpu-vm create tpu-name \ --zone zone \ --accelerator-type accelerator-type \ --version version
Verbindung zur Google Cloud TPU-VM herstellen
Stellen Sie eine SSH-Verbindung zu Ihrer TPU-VM her, indem Sie den folgenden Befehl verwenden:
$ gcloud compute tpus tpu-vm ssh tpu-name --zone zone
Wenn Sie bei der VM angemeldet sind, ändert sich die Shell-Eingabeaufforderung von username@projectname zu username@vm-name:
Pax auf der Google Cloud TPU-VM installieren
Installieren Sie Pax, JAX und libtpu mit den folgenden Befehlen auf Ihrer TPU-VM:
(vm)$ python3 -m pip install -U pip \ python3 -m pip install paxml jax[tpu] -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
Systemüberprüfung
Testen Sie, ob alles korrekt installiert ist. Prüfen Sie dazu, ob JavaX die TPU-Kerne erkennt:
(vm)$ python3 -c "import jax; print(jax.device_count())"
Die Anzahl der TPU-Kerne wird angezeigt. Sie sollten 8 bei Verwendung von v2-8 oder v3-8 bzw. 4 bei einem v4-8 angeben.
Pax-Code auf einer TPU-VM ausführen
Sie können jetzt einen beliebigen Pax-Code ausführen. Die Beispiele von lm_cloud sind ein guter Ausgangspunkt, um Modelle in Pax auszuführen. Mit den folgenden Befehlen wird beispielsweise ein SPMD-Sprachmodell, das auf einem 2B-Parameter-Transformer basiert, anhand synthetischer Daten trainiert.
Die folgenden Befehle zeigen die Trainingsausgabe für ein SPMD-Sprachmodell. Das Training dauert etwa 20 Minuten.
(vm)$ python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=job_log_dir
Im v4-8-Segment sollte die Ausgabe Folgendes enthalten:
Niederlagen und Schrittzeiten
Zusammenfassungs-Tensor bei Schritt=step_# loss = loss
Zusammenfassungs-Tensor bei Schritt=step_# Schritte/Sek. x
Bereinigen
Damit Ihrem Google Cloud-Konto die in dieser Anleitung verwendeten Ressourcen nicht in Rechnung gestellt werden, löschen Sie entweder das Projekt, das die Ressourcen enthält, oder Sie behalten das Projekt und löschen die einzelnen Ressourcen.
Wenn Sie mit Ihrer TPU-VM fertig sind, führen Sie die folgenden Schritte aus, um Ihre Ressourcen zu bereinigen.
Trennen Sie die Verbindung zur Compute Engine-Instanz, falls noch nicht geschehen:
(vm)$ exit
Löschen Sie Ihre Cloud TPU.
$ gcloud compute tpus tpu-vm delete tpu-name --zone zone
Nächste Schritte
Weitere Informationen zu Cloud TPU finden Sie unter: