Dieses Dokument bietet 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 ML-Experimente auf JAX durchführen. Pax konzentriert sich zur Vereinfachung von ML im großen Maßstab durch gemeinsame Nutzung von Infrastrukturkomponenten vorhandenen ML-Frameworks nutzen und 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 zu konfigurieren
Ihr Cloud TPU-Projekt und installieren Sie die für das Training
ein Modell, 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 Google Cloud CLI-Produkten und -Diensten. Wenn Sie die Installieren Sie sie jetzt mithilfe der Anleitung unter Google Cloud CLI installieren.
gcloud-Befehl konfigurieren
Führen Sie 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 sie auch auf der
Google Cloud Console.
$ 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 Ihr Modell und Ihr Code direkt auf der TPU-VM ausgeführt. Sie verbinden sich direkt per SSH mit der TPU-VM. Sie können beliebige Ausführungen Pakete installieren, Logs ansehen und Code zur Fehlerbehebung direkt auf der TPU-VM ausführen.
Erstellen Sie Ihre TPU-VM, indem Sie den folgenden Befehl aus einem in Cloud Shell oder auf Ihrem Computerterminal, Die Google Cloud CLI ist installiert.
Legen Sie die zone je nach Verfügbarkeit in Ihrem Vertrag fest.
siehe TPU-Regionen und -Zonen
wenn nötig.
Legen Sie die Variable accelerator-type auf v2-8, v3-8 oder v4-8 fest.
Variable version auf tpu-vm-base festlegen
für v2- und v3-TPU-Versionen oder tpu-vm-v4-base für v4-TPUs.
$ gcloud compute tpus tpu-vm create tpu-name \
--zone zone \
--accelerator-type accelerator-type \
--version version
Verbindung zu Ihrer 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 in 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
Überprüfen Sie, ob alles korrekt installiert ist. JAX sieht die TPU-Kerne:
(vm)$ python3 -c "import jax; print(jax.device_count())"
Die Anzahl der TPU-Kerne wird angezeigt. Sie sollte 8 sein, wenn Sie eine v2-8 oder v3-8 oder 4, wenn Sie eine v4-8 verwenden.
Pax-Code auf einer TPU-VM ausführen
Sie können jetzt jeden beliebigen Pax-Code ausführen. Die lm_cloud Beispiele sind ein guter Ausgangspunkt, um Modelle in Pax auszuführen. Beispiel: Der Parameter folgenden Befehlen trainieren: Transformator-basiertes SPMD-Sprachmodell von 2B-Parametern mit synthetischen Daten.
Die folgenden Befehle zeigen die Trainingsausgabe für eine SPMD-Sprache model. Sie trainiert in etwa 20 Minuten 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
Bei einem v4-8-Slice sollte die Ausgabe Folgendes enthalten:
Verluste 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 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: