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 auf Basis von JAX. Pax konzentriert sich darauf, ML im großen Maßstab zu vereinfachen. Dazu teilt er Infrastrukturkomponenten mit vorhandenen ML-Frameworks und nutzt die Modellierungsbibliothek Praxis 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 erforderlich sind, 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. Falls noch nicht geschehen, installieren Sie ihn jetzt. Folgen Sie dazu der Anleitung unter Google Cloud CLI installieren.
gcloud-Befehl konfigurieren
(Führen Sie den gcloud auth list aus, um die 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 direkte SSH-Verbindung zur TPU-VM her. Sie können beliebigen Code ausführen, Pakete installieren, Logs ansehen und Code zur Fehlerbehebung direkt auf der TPU-VM ausführen.
Erstellen Sie Ihre TPU-VM. Führen Sie dazu den folgenden Befehl in einer Cloud Shell oder auf dem Computerterminal aus, auf dem die Google Cloud CLI installiert ist.
Legen Sie zone entsprechend der Verfügbarkeit in Ihrem Vertrag fest. Sehen Sie bei Bedarf unter TPU-Regionen und -Zonen nach.
Setzen Sie die Variable accelerator-type auf v2-8, v3-8 oder v4-8.
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())"
Es wird die Anzahl der TPU-Kerne angezeigt. Wenn Sie v2-8 oder v3-8 verwenden, sollte diese 8 sein, wenn Sie v4-8 verwenden, 4.
Pax-Code auf einer TPU-VM ausführen
Sie können jetzt jeden beliebigen Pax-Code ausführen. Die Beispiele lm_cloud sind ein guter Ausgangspunkt, um Modelle in Pax auszuführen. Mit den folgenden Befehlen wird beispielsweise ein SPMD-Sprachmodell auf 2B-Parameter-Transformer anhand synthetischer Daten trainiert.
Die folgenden Befehle zeigen die Trainingsausgabe für ein SPMD-Sprachmodell. Das Training dauert etwa 20 Minuten, also 300 Schritte.
(vm)$ python3 .local/lib/python3.8/site-packages/paxml/main.py --exp=tasks.lm.params.lm_cloud.LmCloudSpmd2BLimitSteps --job_log_dir=job_log_dir
LmCloudSpmd2BLimitSteps
Im v4-8-Slice sollte die Ausgabe Folgendes enthalten:
Niederlagen und Schrittzeiten
Zusammenfassungstensor bei Schritt=step_# loss = loss
Zusammenfassungstensor bei Schritt=step_# Schritte pro Sekunde 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 nicht bereits 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: