Cloud TPU Multislice – Übersicht

Wichtig :Multi-Slice wird nur von Cloud TPU v4 und höheren TPU-Versionen unterstützt.

Cloud TPU Multislice ist eine Full-Stack-Technologie zur Leistungsskalierung, mit der bei einem Trainingsjob mehrere TPU-Segmente innerhalb eines einzelnen Pods oder auf Segmenten in mehreren Pods mit einfacher Datenparallelität verwendet werden können. Bei TPU v4-Chips können für Trainingsjobs also mehr als 4.096 Chips in einem einzigen Durchlauf verwendet werden. Bei Trainingsjobs, für die weniger als 4.096 Chips erforderlich sind, kann ein einzelnes Slice die beste Leistung bieten. Mehrere kleinere Scheiben sind jedoch leichter verfügbar, was zu einer kürzeren Startzeit führt, wenn Multislice mit kleineren Scheiben verwendet wird.

Bei der Bereitstellung in Multislice-Konfigurationen kommunizieren die TPU-Chips in jedem Slice über Inter-Chip-Interconnects (ICI). TPU-Chips in verschiedenen Slices kommunizieren, indem sie Daten an CPUs (Hosts) übertragen, die die Daten wiederum über das Rechenzentrumsnetzwerk (DCN) weiterleiten.

Entwickler müssen keinen Code schreiben, um die DCN-Kommunikation zwischen Slices zu implementieren. Der XLA-Compiler generiert diesen Code für Sie und überschneidet die Kommunikation mit der Berechnung, um die Leistung zu maximieren.

Konzepte

Beschleunigertyp

Die Form jedes TPU-Slices, aus dem ein Multislice besteht. Jedes Slice in einer Anfrage mit mehreren Slices hat denselben Beschleunigertyp. Ein Beschleunigertyp besteht aus einem TPU-Typ (v4 oder v5e) gefolgt von der Anzahl der Tensorkerne. v4-128 gibt beispielsweise eine TPU v4 mit 128 TensorCores an.

Automatische Reparatur

Wenn bei einem Slice ein Wartungsereignis, eine Voraktivierung oder ein Hardwarefehler auftritt, erstellt Cloud TPU einen neuen Slice. In dem seltenen Fall, dass nicht genügend Ressourcen für die Erstellung eines neuen Slices vorhanden sind, wird die Erstellung erst abgeschlossen, wenn Hardware verfügbar ist. Nachdem der neue Sliver erstellt wurde, werden alle anderen Sliver in der Multislice-Umgebung neu gestartet, damit das Training fortgesetzt werden kann.Mit einem ordnungsgemäß konfigurierten Startskript kann das Trainingsskript automatisch ohne Nutzereingriff neu gestartet werden, indem es vom letzten Checkpoint geladen und fortgesetzt wird.

Dataset

Daten, die von einem Modell für das Training oder die Inferenz verwendet werden.

Rechenzentrumsnetzwerk (DCN)

Ein Netzwerk mit höherer Latenz und geringerem Durchsatz (im Vergleich zu ICI), das TPU-Slices in einer Multislice-Konfiguration verbindet.

Parallele Ausführung

Wenn alle TPU-Scheiben gleichzeitig bereitgestellt werden, wird sichergestellt, dass entweder alle oder keine der Scheiben bereitgestellt werden.

Moderator:in

Ein Host ist ein physischer Computer, auf dem VMs ausgeführt werden. Auf einem Host können maximal vier VMs gleichzeitig ausgeführt werden. Jede VM hat eine dedizierte TPU.

Inferenz

Vortrainiertes Machine-Learning-Modell auf einen Host laden und Vorhersagen auf Daten treffen.

Interchip-Interconnect (ICI)

Interne Verbindungen mit hoher Geschwindigkeit und niedriger Latenz, die TPUs innerhalb eines TPU-Pods verbinden.

Mehrfachschicht

Zwei oder mehr TPU-Chip-Slices, die über das DCN kommunizieren können.

Knoten

Im Multislice-Kontext bezieht sich „Knoten“ auf ein einzelnes TPU-Stück. Jedem TPU-Speicherbereich in einem Multi-Slice wird eine Knoten-ID zugewiesen.

Pod

Eine Sammlung von TPU-Chips, die über dedizierte ICI-Netzwerkschnittstellen verbunden sind. Mit einem Pod können Sie die Verarbeitungslast auf mehrere TPUs verteilen.

In der Warteschlange befindliche Ressource (QR)

Eine Darstellung von TPU-Ressourcen, die zum Einreihen und Verwalten einer Anfrage für eine TPU-Umgebung mit einer oder mehreren Slices verwendet wird.

Startskript

Ein standardmäßiges Compute Engine-Startskript, das jedes Mal ausgeführt wird, wenn eine VM gestartet oder neu gestartet wird. Bei Multislice wird sie in der Anfrage zum Erstellen des QR-Codes angegeben. Weitere Informationen zu Cloud TPU-Startscripts finden Sie unter TPU-Ressourcen verwalten.

TPU-Speichereinheit

Logischer Teilbereich eines TPU-Pod, der aus TPU-Chips besteht. Alle Chips in einem Slice kommunizieren über das ICI-Netzwerk miteinander.

TPU-VM

Eine virtuelle Maschine mit Linux, die Zugriff auf die zugrunde liegenden TPUs hat. Bei TPUs der Version 4 hat jede TPU-VM direkten Zugriff auf vier Chips. Manchmal wird eine TPU-VM auch als Worker bezeichnet.

Tensor

Eine Datenstruktur, die verwendet wird, um mehrdimensionale Daten in einem Modell für maschinelles Lernen darzustellen.

Tensor Processing Unit (TPU)

Der intern von Google entwickelte Chip zur Beschleunigung von ML. Sie sind für schnelle und energieeffiziente Berechnungen bei wichtigen Aufgaben des maschinellen Lernens wie der Matrixmultiplikation konzipiert.

Arten von Cloud TPU-Kapazitäten

TPUs können aus verschiedenen Arten von Kapazität erstellt werden (siehe „Nutzungsoptionen“ im Hilfeartikel So funktionieren die TPU-Preise) :

• Reservierung: Wenn Sie eine Reservierung nutzen möchten, müssen Sie eine Reservierungsvereinbarung mit Google haben. Verwenden Sie das Flag --reserved, wenn Sie Ressourcen erstellen.
• Spot: Das Kontingent auf Abruf wird mit Spot-VMs erreicht. Ihre Ressourcen werden möglicherweise vorzeitig freigegeben, um Platz für Anfragen für einen Job mit höherer Priorität zu schaffen. Verwenden Sie beim Erstellen Ihrer Ressourcen das Flag --spot.
• On-Demand: Zielt auf das On-Demand-Kontingent ab, für das keine Reservierung erforderlich ist und das nicht vorrangig verwendet wird. Die TPU-Anfrage wird in eine On-Demand-Kontingentwarteschlange gestellt, die von Cloud TPU angeboten wird. Die Verfügbarkeit von Ressourcen ist nicht garantiert. Standardmäßig ausgewählt, keine Flags erforderlich.

Hinweis: Cloud TPU schränkt die Größe der Sliver für alle Beschleunigertypen ein. Informationen zu den Limits für die Größe von Segmenten bei On-Demand-Anfragen finden Sie unter Standardsegmentgrößen für On-Demand-QR-Codes.

Mehr erfahren

Wenn Sie noch keine TPUs verwendet haben, installieren Sie zuerst die Google Cloud CLI und richten Sie Ihre Cloud TPU-Umgebung ein. Wenn Sie Multislice verwenden möchten, müssen Ihre TPU-Ressourcen als in die Warteschlange gestellte Ressourcen verwaltet werden.

Wenn Sie bereits TPU v4 verwenden und eine Reservierung haben, müssen Sie diese möglicherweise in ein neues Reservierungssystem migrieren. Weitere Informationen erhalten Sie von Ihrem Google Cloud-Kundenbetreuer.

Beispiel für den Einstieg

In dieser Anleitung wird Code aus dem GitHub-Repository von MaxText verwendet. MaxText ist eine leistungsstarke, beliebig skalierbare, Open-Source- und gut getestete grundlegende LLM, die in Python und Jax geschrieben wurde. MaxText wurde entwickelt, um effizient auf Cloud TPU zu trainieren.

Der Code in shardings.py soll Ihnen den Einstieg in verschiedene Parallelisierungsoptionen erleichtern. Dazu gehören beispielsweise Datenparallelität, vollständig fragmentierte Datenparallelität (FSDP) und Tensorparallelität. Der Code kann von einer einzelnen Schleife auf Multislice-Umgebungen skaliert werden.

ICI-Parallelität

ICI bezieht sich auf die Hochgeschwindigkeits-Interconnect-Verbindung, die die TPUs in einem einzelnen Slice verbindet. ICI-Sharding entspricht dem Sharding innerhalb eines Segments. shardings.py bietet drei ICI-Parallelisierungsparameter:

• ici_data_parallelism
• ici_fsdp_parallelism
• ici_tensor_parallelism

Die Werte, die Sie für diese Parameter angeben, bestimmen die Anzahl der Shards für jede Parallelisierungsmethode.

Diese Eingaben müssen so eingeschränkt sein, dass ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism der Anzahl der Chips im Slice entspricht.

Die folgende Tabelle enthält Beispielnutzereingaben für die ICI-Parallelität der vier Chips, die in Version 4–8 verfügbar sind:

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
4-Wege-FSDP	1	4	1
4-Wege-Tensor-Parallelität	1	1	4
Zwei-Wege-FSDP + Zwei-Wege-Tensor-Parallelität	1	2	2

Hinweis: ici_data_parallelism sollte in den meisten Fällen bei 1 belassen werden, da das ICI-Netzwerk schnell genug ist, um fast immer FSDP dem Datenparallelismus vorzuziehen.

In diesem Beispiel wird davon ausgegangen, dass Sie mit dem Ausführen von Code auf einem einzelnen TPU-Speicherbereich vertraut sind, z. B. wie im Artikel Berechnung mit JAX auf einer Cloud TPU-VM ausführen. In diesem Beispiel wird gezeigt, wie shardings.py auf einen einzelnen Ausschnitt angewendet wird.

1. Richten Sie die Umgebung ein:

   $ gcloud auth login
   $ gcloud config set project your-project-id
   $ gcloud config set compute/zone your-zone

2. Erstellen Sie SSH-Schlüssel für gcloud. Wir empfehlen, das Passwort leer zu lassen. Drücken Sie dazu nach Ausführung des folgenden Befehls zweimal die Eingabetaste. Wenn Sie aufgefordert werden, die Datei google_compute_engine zu ersetzen, weil sie bereits vorhanden ist, ersetzen Sie die vorhandene Version.

   $ ssh-keygen -f ~/.ssh/google_compute_engine

3. Stellen Sie Ihre TPUs bereit:

   gcloud

   $ gcloud compute tpus queued-resources \
       create YOUR_QR_ID \
       --accelerator-type your-accelerator-type \
       --runtime-version tpu-ubuntu2204-base \
       --node-id qr-id \
       [--reserved |--spot]

   Beschreibung der Befehls-Flags

   YOUR_QR_ID

   Ein benutzerdefinierter String, der die QR-Anfrage identifiziert.

   accelerator-type

   Mit dem Beschleunigertyp geben Sie die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Weitere Informationen zu den unterstützten Beschleunigertypen für jede TPU-Version finden Sie unter TPU-Versionen.

   runtime-version

   Die Cloud TPU-Softwareversion.

   node-id

   Die ID der TPU-Ressourcen, die als Reaktion auf die QR-Anfrage erstellt werden.

   reserved

   Verwenden Sie beim Erstellen der Slices eine Reservierung.

   spot

   Verwenden Sie beim Erstellen der Slices Spot-VMs.

   Die Google Cloud CLI unterstützt nicht alle Optionen zum Erstellen von QR-Codes, z. B. Tags. Weitere Informationen finden Sie unter QR-Codes erstellen.

   Console

   1. Rufen Sie in der Google Cloud Console die Seite TPUs auf:

      TPUs aufrufen

   2. Klicken Sie auf TPU erstellen.

   3. Geben Sie im Feld Name einen Namen für die TPU ein.

   4. Wählen Sie im Feld Zone die Zone aus, in der die TPU erstellt werden soll.

   5. Wählen Sie im Feld TPU-Typ einen Beschleunigertyp aus. Mit dem Beschleunigertyp geben Sie die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Weitere Informationen zu den unterstützten Beschleunigertypen für jede TPU-Version finden Sie unter TPU-Versionen.

   6. Wählen Sie im Feld TPU-Softwareversion eine Softwareversion aus. Beim Erstellen einer Cloud TPU-VM gibt die TPU-Softwareversion die Version der zu installierenden TPU-Laufzeit an. Weitere Informationen finden Sie unter TPU-VM-Images.

   7. Klicken Sie auf den Schalter Warteschlange aktivieren.

   8. Geben Sie im Feld Name der anstehenden Ressource einen Namen für die anstehende Ressourcenanfrage ein.

   9. Klicken Sie auf Erstellen, um die anstehende Ressourcenanfrage zu erstellen.

4. Warten Sie, bis die in der Warteschlange befindliche Ressource den Status ACTIVE hat. Das bedeutet, dass sich die Worker-Knoten im Status READY befinden. Sobald die Bereitstellung der Ressourcen in der Warteschlange gestartet wurde, kann es je nach Größe der Ressourcen ein bis fünf Minuten dauern, bis sie abgeschlossen ist. Sie können den Status einer anstehenden Ressourcenanfrage mit der gcloud CLI oder der Google Cloud Console prüfen:

   gcloud

   $ gcloud compute tpus queued-resources \
       list --filter=YOUR_QR_ID

   Console

   1. Rufen Sie in der Google Cloud Console die Seite TPUs auf:

      TPUs aufrufen

   2. Klicken Sie auf den Tab In der Warteschlange befindliche Ressourcen.

   3. Klicken Sie auf den Namen Ihrer in der Warteschlange befindlichen Ressourcenanfrage.

5. Ein v4-8-Slice hat eine einzelne TPU-VM. Stellen Sie eine SSH-Verbindung zur TPU-VM her:

   $ gcloud compute tpus tpu-vm ssh YOUR_QR_ID

6. Klonen Sie MaxText (einschließlich shardings.py) auf Ihre TPU-VM.

7. Führen Sie im MaxText-Repository-Verzeichnis das Einrichtungsskript aus, um JAX und andere Abhängigkeiten auf Ihrem TPU-Speichereinsatz zu installieren. Die Ausführung des Einrichtungsscripts dauert einige Minuten.

   $ bash setup.sh

8. Führen Sie den folgenden Befehl aus, um shardings.py auf Ihrem TPU-Speicherbereich auszuführen.

   $ python3 pedagogical_examples/shardings.py \
     --ici_fsdp_parallelism 4 \
     --batch_size 131072 \
     --embedding_dimension 2048

   Die Ergebnisse finden Sie in den Protokollen. Ihre TPUs sollten etwa 260 TFLOP/Sekunde erreichen, was einer beeindruckenden FLOP-Auslastung von über 90%entspricht. In diesem Fall haben wir ungefähr die maximale Batchgröße ausgewählt, die in den High Bandwidth Memory (HBM) der TPU passt.

9. Sie können auch andere Sharding-Strategien über ICI ausprobieren. Probieren Sie beispielsweise die folgende Kombination aus:

   $ python3 pedagogical_examples/shardings.py \
     --ici_tensor_parallelism 4 \
     --batch_size 131072 \
     --embedding_dimension 2048

10. Löschen Sie die in die Warteschlange gestellte Ressource und das TPU-Stück, wenn Sie fertig sind. Führen Sie diese Schritte zur Bereinigung in der Umgebung aus, in der Sie den Sliver eingerichtet haben. Führen Sie zuerst exit aus, um die SSH-Sitzung zu beenden. Das Löschen dauert zwei bis fünf Minuten. Wenn Sie die gcloud CLI verwenden, können Sie diesen Befehl mit dem optionalen Flag --async im Hintergrund ausführen.

    Wichtig :Ressourcen in der Warteschlange müssen nach der Verwendung gelöscht werden. Andernfalls wird Ihr Kontingent für Warteschlangenressourcen weiter verbraucht. Dieses wird getrennt vom TPU-Kontingent erfasst.

    gcloud

    $ gcloud compute tpus queued-resources
        delete YOUR_QR_ID --force (--async)

    Console

    1. Rufen Sie in der Google Cloud Console die Seite TPUs auf:

       TPUs aufrufen

    2. Klicken Sie auf den Tab In der Warteschlange befindliche Ressourcen.

    3. Klicken Sie das Kästchen neben der anstehenden Ressourcenanfrage an.

    4. Klicken Sie auf Löschen.

Mehrfachscheiben-Sharding mit DCN-Parallelität

Das shardings.py-Script verwendet drei Parameter, die die DCN-Parallelität angeben. Sie entsprechen der Anzahl der Shards der einzelnen Datenparallelitätstypen:

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

Die Werte dieser Parameter müssen so eingeschränkt sein, dass dcn_data_parallelism * dcn_fsdp_parallelism * dcn_tensor_parallelism der Anzahl der Scheiben entspricht.

Verwenden Sie für das Beispiel mit zwei Segmenten --dcn_data_parallelism = 2.

	dcn_data_parallelism	dcn_fsdp_parallelism	dcn_tensor_parallelism	Anzahl der Segmente
Bidirektionale Datenparallelität	2	1	1	2

dcn_tensor_parallelism sollte immer auf 1 festgelegt sein, da das DCN für eine solche Sharding-Technologie nicht geeignet ist. Für typische LLM-Arbeitslasten auf v4-Chips sollte dcn_fsdp_parallelism ebenfalls auf 1 und dcn_data_parallelism auf die Anzahl der Slices festgelegt werden. Dies ist jedoch anwendungsabhängig.

Wenn Sie die Anzahl der Scheiben erhöhen (vorausgesetzt, Sie halten die Scheibengröße und den Batch pro Scheibe konstant), erhöhen Sie die Datenparallelität.

shardings.py in einer Multislice-Umgebung ausführen

Sie können shardings.py in einer Multislice-Umgebung mit multihost_runner.py oder durch Ausführen von shardings.py auf jeder TPU-VM ausführen. Hier verwenden wir multihost_runner.py. Die folgenden Schritte ähneln denen unter Einstieg: Schnelle Tests für mehrere Segmente im MaxText-Repository. Der Unterschied besteht darin, dass hier shardings.py anstelle des komplexeren LLM in train.py verwendet wird.

Das multihost_runner.py-Tool ist für schnelle Tests optimiert und verwendet wiederholt dieselben TPUs. Da das multihost_runner.py-Script auf lang anhaltenden SSH-Verbindungen basiert, empfehlen wir es nicht für lang laufende Jobs. Wenn Sie einen längeren Job ausführen möchten (z. B. Stunden oder Tage), empfehlen wir multihost_job.py.

In dieser Anleitung wird der Begriff Runner für den Computer verwendet, auf dem Sie das multihost_runner.py-Script ausführen. Wir verwenden den Begriff Worker für die TPU-VMs, aus denen Ihre Slices bestehen. Sie können multihost_runner.py auf einem lokalen Computer oder auf einer beliebigen Compute Engine-VM im selben Projekt wie Ihre Slices ausführen. Das Ausführen von multihost_runner.py auf einem Worker wird nicht unterstützt.

multihost_runner.py stellt automatisch eine SSH-Verbindung zu TPU-Workern her.

In diesem Beispiel wird shardings.py über zwei v4-16-Segmente ausgeführt, insgesamt vier VMs und 16 TPU-Chips. Sie können das Beispiel so ändern, dass es auf mehr TPUs ausgeführt wird.

Umgebung einrichten

1. Klonen Sie MaxText auf dem Computer des Runners.

2. Rufen Sie das Repository-Verzeichnis auf.

3. Erstellen Sie SSH-Schlüssel für gcloud. Wir empfehlen, das Passwort leer zu lassen. Drücken Sie nach Ausführung des folgenden Befehls zweimal die Eingabetaste. Wenn Sie aufgefordert werden, dass die Datei google_compute_engine bereits vorhanden ist, wählen Sie aus, dass die vorhandene Version nicht beibehalten werden soll.

     $ ssh-keygen -f ~/.ssh/google_compute_engine
     

4. Fügen Sie eine Umgebungsvariable hinzu, um die Anzahl der TPU-Scheiben auf 2 festzulegen.

     $ export SLICE_COUNT=2
     

5. Erstellen Sie eine Multislice-Umgebung mit dem Befehl queued-resources create oder der Google Cloud Console.

   gcloud

   Hinweis: Sie müssen das entsprechende Kontingent haben, bevor Sie --reserved, --spot oder das Standardkontingent für On-Demand-Aktionen auswählen. Informationen zu Kontingenttypen finden Sie in der Kontingentrichtlinie.

   Der folgende Befehl zeigt, wie Sie eine v4-Multislice-TPU erstellen. Wenn Sie v5e verwenden möchten, geben Sie eine v5e-accelerator-type (z. B. v5litepod-16) und die v5e-runtime-version (v2-alpha-tpuv5-lite) an.

   $ gcloud compute tpus queued-resources \
       create YOUR_QR_ID \
       --accelerator-type=your-accelerator-type \
       --runtime-version=tpu-vm-runtime-version \
       --node-count=node-count \
       --node-prefix=YOUR_QR_ID \
       [--reserved|--spot]

   Beschreibung der Befehls-Flags

   YOUR_QR_ID

   Ein benutzerdefinierter String, der die QR-Anfrage identifiziert.

   accelerator-type

   Mit dem Beschleunigertyp geben Sie die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Multi-Slice wird nur von Cloud TPU v4 und höheren TPU-Versionen unterstützt. Weitere Informationen zu den unterstützten Beschleunigertypen für jede TPU-Version finden Sie unter TPU-Versionen.

   runtime-version

   Die Cloud TPU-Softwareversion.

   node-count

   Die Anzahl der zu erstellenden Scheiben.

   node-prefix

   Das Präfix, mit dem Namen für die einzelnen Segmente generiert werden. Dem Präfix wird für jeden Ausschnitt eine Zahl angehängt. Wenn Sie beispielsweise node-prefix auf mySlice festlegen, werden die Slices mySlice-0, mySlice-1 usw. benannt.

   reserved

   Verwenden Sie beim Erstellen der Slices eine Reservierung.

   spot

   Verwenden Sie beim Erstellen der Slices Spot-VMs.

   Console

   1. Rufen Sie in der Google Cloud Console die Seite TPUs auf:

      TPUs aufrufen

   2. Klicken Sie auf TPU erstellen.

   3. Geben Sie im Feld Name einen Namen für die TPU ein.

   4. Wählen Sie im Feld Zone die Zone aus, in der die TPU erstellt werden soll.

   5. Wählen Sie im Feld TPU-Typ einen Beschleunigertyp aus. Mit dem Beschleunigertyp geben Sie die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Multi-Slice wird nur von Cloud TPU v4 und höheren TPU-Versionen unterstützt. Weitere Informationen zu den unterstützten Beschleunigertypen für jede TPU-Version finden Sie unter TPU-Versionen.

   6. Wählen Sie im Feld TPU-Softwareversion eine Softwareversion aus. Beim Erstellen einer Cloud TPU-VM gibt die TPU-Softwareversion die Version der zu installierenden TPU-Laufzeit an. Weitere Informationen finden Sie unter TPU-VM-Images.

   7. Klicken Sie auf den Schalter Warteschlange aktivieren.

   8. Geben Sie im Feld Name der anstehenden Ressource einen Namen für die anstehende Ressourcenanfrage ein.

   9. Klicken Sie das Kästchen Multislice-TPU erstellen an.

   10. Geben Sie im Feld Anzahl der Scheiben die Anzahl der Scheiben ein, die Sie erstellen möchten.

   11. Klicken Sie auf Erstellen, um die anstehende Ressourcenanfrage zu erstellen.

6. Wenn die Bereitstellung der Ressourcen in der Warteschlange beginnt, kann es je nach Größe der Ressourcen bis zu fünf Minuten dauern, bis sie abgeschlossen ist. Warten Sie, bis die in der Warteschlange befindliche Ressource den Status ACTIVE hat. Sie können den Status einer anstehenden Ressourcenanfrage mit der gcloud CLI oder der Google Cloud Console prüfen:

   gcloud

   $ gcloud compute tpus queued-resources list \
       --filter=YOUR_QR_ID

   Die Ausgabe sollte in etwa so aussehen:

   NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
   ...
   que-res-id  us-central2-b  4           v4-16             ACTIVE
   ...

   Console

   1. Rufen Sie in der Google Cloud Console die Seite TPUs auf:

      TPUs aufrufen

   2. Klicken Sie auf den Tab In der Warteschlange befindliche Ressourcen.

   3. Klicken Sie auf den Namen Ihrer in der Warteschlange befindlichen Ressourcenanfrage.

   Wenden Sie sich an Ihren Google Cloud-Kundenbetreuer, wenn der QR-Status länger als 15 Minuten den Status WAITING_FOR_RESOURCES oder PROVISIONING hat.

7. Installieren Sie die Abhängigkeiten:

   $ python3 multihost_runner.py \
       --TPU_PREFIX=YOUR_QR_ID \
       --COMMAND="bash setup.sh"

8. Führen Sie shardings.py mit multihost_runner.py auf jedem Worker aus.

   $ python3 multihost_runner.py \
       --TPU_PREFIX=YOUR_QR_ID \
       --COMMAND="python3 pedagogical_examples/shardings.py \
       --dcn_data_parallelism $SLICE_COUNT \
       --ici_fsdp_parallelism 8 \
       --batch_size 131072 \
       --embedding_dimension 2048"

   In den Protokolldateien wird eine Leistung von etwa 230 TFLOP/s angezeigt.

9. Bereinigen Sie die TPUs und die in die Warteschlange gestellte Ressource, wenn Sie fertig sind. Das Löschen dauert zwei bis fünf Minuten. Wenn Sie die gcloud CLI verwenden, können Sie diesen Befehl mit dem optionalen Flag --async im Hintergrund ausführen.

Arbeitslast auf Multislice skalieren

Bevor Sie Ihr Modell in einer Multislice-Umgebung ausführen, nehmen Sie die folgenden Codeänderungen vor:

• Verwenden Sie beim Erstellen des Mesh jax.experimental.mesh_utils.create_hybrid_device_mesh anstelle von jax.experimental.mesh_utils.create_device_mesh.

Das sollten die einzigen Codeänderungen sein, die beim Wechsel zu Multislice erforderlich sind. Um eine hohe Leistung zu erzielen, muss das DCN auf datenparrallele, vollständig shardete datenparrallele oder pipelineparallele Achsen abgebildet werden. Leistungsaspekte und Sharding-Strategien werden im Artikel Sharding mit Multislice für maximale Leistung ausführlicher behandelt.

Um zu prüfen, ob dein Code auf alle Geräte zugreifen kann, kannst du prüfen, ob len(jax.devices()) der Anzahl der Chips in deiner Multislice-Umgebung entspricht. Wenn Sie beispielsweise vier Scheiben von v4-16 verwenden, haben Sie acht Chips pro Scheibe × 4 Scheiben, sodass len(jax.devices()) den Wert 32 zurückgeben sollte.

Scheibengrößen für Multi-Scheiben-Umgebungen auswählen

Wenn Sie eine lineare Beschleunigung erzielen möchten, fügen Sie neue Segmente mit derselben Größe wie das vorhandene Segment hinzu. Wenn Sie beispielsweise einen v4-512-Speicherplatz verwenden, kann Multislice die Leistung ungefähr verdoppeln, indem ein zweiter v4-512-Speicherplatz hinzugefügt und die globale Batchgröße verdoppelt wird. Weitere Informationen finden Sie unter Sharding mit Multislice für maximale Leistung.

Job auf mehreren Slices ausführen

Es gibt drei verschiedene Ansätze, um Ihre benutzerdefinierte Arbeitslast in einer Multi-Slice-Umgebung auszuführen:

1. Mit dem Test-Runner-Script multihost_runner.py
2. Mit dem Produktions-Runner-Script multihost_job.py
3. Manueller Ansatz

Script für Testausführung

Das Script multihost_runner.py verteilt Code in einer vorhandenen Multislice-Umgebung, führt den Befehl auf jedem Host aus, kopiert die Protokolle zurück und überwacht den Fehlerstatus jedes Befehls. Das multihost_runner.py-Script ist in der MaxText-README dokumentiert.

Da multihost_runner.py persistente SSH-Verbindungen aufrechterhält, eignet es sich nur für relativ kleine, relativ kurz laufende Tests. Sie können die Schritte im Tutotial für multihost_runner.py an Ihre Arbeitslast und Hardwarekonfiguration anpassen.

Script für den Produktionsläufer

Für Produktionsjobs, die für Ausfallsicherheit bei Hardwarefehlern und anderen Vorwegnahmen sorgen müssen, ist es am besten, sie direkt in die Create Queued Resource API einzubinden. Als Beispiel dient multihost_job.py, mit dem der API-Aufruf „Created Queued Resource“ mit dem entsprechenden Startskript ausgelöst wird, um das Training auszuführen und bei Vorrang fortzusetzen. Das multihost_job.py-Script ist in der MaxText-README dokumentiert.

Da für multihost_job.py für jeden Durchlauf Ressourcen bereitgestellt werden müssen, ist der Iterationszyklus nicht so schnell wie bei multihost_runner.py.

Manueller Ansatz

Wir empfehlen, multihost_runner.py oder multihost_job.py zu verwenden oder anzupassen, um Ihre benutzerdefinierte Arbeitslast in Ihrer Multislice-Konfiguration auszuführen. Wenn Sie Ihre Umgebung jedoch lieber direkt über QR-Befehle bereitstellen und verwalten möchten, lesen Sie den Hilfeartikel Multi-Slice-Umgebung verwalten.

Multislice-Umgebung verwalten

Wenn Sie QR-Codes manuell bereitstellen und verwalten möchten, ohne die im MaxText-Repository bereitgestellten Tools zu verwenden, lesen Sie die folgenden Abschnitte.

In die Warteschlange gestellte Ressourcen erstellen

gcloud

1. Legen Sie die folgenden Umgebungsvariablen fest, bevor Sie Kapazitäten bereitstellen:

   $ export YOUR_QR_ID=your-queued-resource-id
   $ export PROJECT=your-project-name
   $ export ZONE=us-central2-b
   $ export NETWORK_NAME=your-network-name
   $ export SUBNETWORK_NAME=your-subnetwork-name
   $ export RUNTIME_VERSION=tpu-ubuntu2204-base
   $ export ACCELERATOR_TYPE=v4-16
   $ export SLICE_COUNT=4
   $ export STARTUP_SCRIPT="#!/bin/bash\n ..."
   $ gcloud config set project project-name
   $ gcloud config set compute/zone zone

   Variablenbeschreibungen

   Eingabe	Beschreibung
   YOUR_QR_ID	Die vom Nutzer zugewiesene ID der in der Warteschlange befindlichen Ressource.
   PROJEKT	Name des Google Cloud-Projekts
   ZONE	Gibt die Zone an, in der die Ressourcen erstellt werden sollen.
   NETWORK_NAME	Name der VPC-Netzwerke.
   SUBNETWORK_NAME	Name des Subnetzes in VPC-Netzwerken
   RUNTIME_VERSION	Die Cloud TPU-Softwareversion.
   ACCELERATOR_TYPE	v4-16
   Bsp_TAG_1, Bsp_TAG_2…	Tags, mit denen gültige Quellen oder Ziele für Netzwerk-Firewalls angegeben werden
   SLICE_COUNT	Anzahl der Segmente. Maximal 256 Scheiben.
   STARTUP_SCRIPT	Wenn Sie ein Startskript angeben, wird es ausgeführt, wenn das TPU-Stück bereitgestellt oder neu gestartet wird.

2. Erstellen Sie mit dem folgenden Befehl eine Ressourcenanfrage in der Warteschlange:

   $ gcloud compute tpus queued-resources \
       create ${YOUR_QR_ID} \
       --project your-project-id \
       --zone your-zone \
       --node-count ${SLICE_COUNT} \
       --accelerator-type ${ACCELERATOR_TYPE} \
       --runtime-version ${RUNTIME_VERSION} \
       --network ${NETWORK_NAME} \
       --subnetwork ${SUBNETWORK_NAME} \
       --tags ${EXAMPLE_TAG_1},${EXAMPLE_TAG_2} \ --metadata=startup-script='${STARTUP_SCRIPT}'
       [--reserved|--spot]

   Beschreibung der Befehls-Flags

   YOUR_QR_ID

   Ein benutzerdefinierter String, der die anstehende Ressourcenanfrage identifiziert.

   project

   Das Google Cloud-Projekt, in dem Sie die anstehende Ressourcenanfrage erstellen.

   zone

   Die Google Cloud-Zone, in der die Ressourcen in der Warteschlange erstellt werden sollen.

   node-count

   Die Anzahl der zu erstellenden Scheiben.

   accelerator-type

   Mit dem Beschleunigertyp geben Sie die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Multi-Slice wird nur von Cloud TPU v4 und höheren TPU-Versionen unterstützt. Weitere Informationen zu den unterstützten Beschleunigertypen für jede TPU-Version finden Sie unter TPU-Versionen.

   runtime-version

   Die Cloud TPU-Softwareversion.

   network

   Der Name eines VPC-Netzwerk, an das die TPU-Ressource angehängt werden soll.

   subnetwork

   Der Name eines VPC-Subnetzes, an das die TPU-Ressource angehängt werden soll.

   reserved

   Verwenden Sie beim Erstellen der Slices eine Reservierung.

   spot

   Verwenden Sie beim Erstellen der Slices Spot-VMs.

Achten Sie darauf, dass Sie das entsprechende Kontingent haben, bevor Sie --reserved, --spot oder das Standardkontingent für On-Demand-Anzeigeaufträge auswählen. Informationen zu Kontingenttypen finden Sie in der Kontingentrichtlinie.

curl

1. Erstellen Sie eine Datei namens queued-resource-req.json und kopieren Sie die folgende JSON-Datei hinein.

   {
   "guaranteed": { "reserved": true },
   "tpu": {
       "node_spec": [
       {
       "parent": "projects/your-project-number/locations/your-zone",
           "node": {
           "accelerator_type": "accelerator-type",
           "runtime_version": "tpu-vm-runtime-version",
           "network_config": {
               "network": "your-network-name",
               "subnetwork": "your-subnetwork-name",
               "enable_external_ips": true
           },
           "tags" : ["example-tag-1"]
           "metadata": {
               "startup-script": "your-startup-script"
           }
       },
       "multi_node_params": {
           "node_count": slice-count,
           "node_id_prefix": "your-queued-resource-id"
       }
       }
       ]
   }
   }

   Ersetzen Sie die folgenden Werte:

   • your-project-number: Ihre Google Cloud-Projektnummer
   • your-zone – die Zone, in der Sie die Ressourcen in der Warteschlange erstellen möchten
   • accelerator-type: Version und Größe eines einzelnen Slices. Multi-Slice wird nur von Cloud TPU v4 und höheren TPU-Versionen unterstützt.
   • tpu-vm-runtime-version: Die zu verwendende TPU-VM-Laufzeitversion.
   • your-network-name – Optional: ein Netzwerk, an das die anstehende Ressource angehängt wird
   • your-subnetwork-name – Optional: ein Subnetzwerk, an das die in der Warteschlange befindliche Ressource angehängt werden soll
   • example-tag-1 – Optionaler, beliebiger Tag-String
   • your-startup-script: Startskript, das ausgeführt wird, wenn die in der Warteschlange befindliche Ressource zugewiesen wird
   • slice-count – Die Anzahl der TPU-Slices in Ihrer Multislice-Umgebung
   • YOUR_QR_ID: Die vom Nutzer angegebene ID für die in der Warteschlange befindliche Ressource

   Weitere Informationen zu allen verfügbaren Optionen findest du in der Dokumentation zur REST Queued Resource API.

   Hinweis: node_spec ist zwar eine Liste, aber es wird nur ein node_spec-Eintrag unterstützt.

   Wenn Sie die Spot-Kapazität verwenden möchten, ersetzen Sie Folgendes:

   "guaranteed": { "reserved": true } mit "spot": {}

   Entfernen Sie die Zeile, um die Standardkapazität für On-Demand-Kapazität zu verwenden.

2. Reiche die Anfrage zum Erstellen der anstehenden Ressource mit der JSON-Nutzlast ein:

   $ curl -X POST -H "Authorization: Bearer $(gcloud auth print-access-token)" \
   -H "Content-Type: application/json" \
   -d @queuedresourcereq.json \
   https://tpu.googleapis.com/v2alpha1/projects/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=YOUR_QR_ID

   Ersetzen Sie die folgenden Werte:

   • your-project-id: Ihre Google Cloud-Projekt-ID
   • your-zone – die Zone, in der Sie die Ressourcen in der Warteschlange erstellen möchten
   • YOUR_QR_ID: Die vom Nutzer angegebene ID für die in der Warteschlange befindliche Ressource

Die Antwort sollte in etwa so aussehen:

{
"name": "projects/<your-project-id>/locations/<your-zone>/operations/operation-<your-qr-guid>",
"metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-11-01T00:17:05.742546311Z",
    "target": "projects/<your-project-id>/locations/<your-zone>/queuedResources/<your-qa-id>",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
},
"done": false
}

Verwenden Sie den GUID-Wert am Ende des Stringwerts für das name-Attribut, um Informationen zur anstehenden Ressourcenanfrage zu erhalten.

Console

1. Rufen Sie in der Google Cloud Console die Seite TPUs auf:

   TPUs aufrufen

2. Klicken Sie auf TPU erstellen.

3. Geben Sie im Feld Name einen Namen für die TPU ein.

4. Wählen Sie im Feld Zone die Zone aus, in der die TPU erstellt werden soll.

5. Wählen Sie im Feld TPU-Typ einen Beschleunigertyp aus. Mit dem Beschleunigertyp geben Sie die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Multi-Slice wird nur von Cloud TPU v4 und höheren TPU-Versionen unterstützt. Weitere Informationen zu den unterstützten Beschleunigertypen für jede TPU-Version finden Sie unter TPU-Versionen.

6. Wählen Sie im Feld TPU-Softwareversion eine Softwareversion aus. Beim Erstellen einer Cloud TPU-VM gibt die TPU-Softwareversion die Version der zu installierenden TPU-Laufzeit an. Weitere Informationen finden Sie unter TPU-VM-Images.

7. Klicken Sie auf den Schalter Warteschlange aktivieren.

8. Geben Sie im Feld Name der anstehenden Ressource einen Namen für die anstehende Ressourcenanfrage ein.

9. Klicken Sie das Kästchen Multislice-TPU erstellen an.

10. Geben Sie im Feld Anzahl der Scheiben die Anzahl der Scheiben ein, die Sie erstellen möchten.

11. Klicken Sie auf Erstellen, um die anstehende Ressourcenanfrage zu erstellen.

Status einer in die Warteschlange gestellten Ressource abrufen

gcloud

$ gcloud compute tpus queued-resources describe ${YOUR_QR_ID}

Bei einer Ressourcenwarteschlange im Status ACTIVE sieht die Ausgabe so aus:

...
state:
    state: ACTIVE
...

curl

$ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2/projects/your-project-id/locations/your-zone/queuedResources/${YOUR_QR_ID}

Bei einer Ressourcenwarteschlange im Status ACTIVE sieht die Ausgabe so aus:

{
"name": your-queued-res,
"tpu": {
    "nodeSpec": [
    {
        ... // node 1
    },
    {
        ... // node 2
    },
    ...
    ]
},
...
"state": "ACTIVE"
}

Console

1. Rufen Sie in der Google Cloud Console die Seite TPUs auf:

   TPUs aufrufen

2. Klicken Sie auf den Tab In der Warteschlange befindliche Ressourcen.

3. Klicken Sie auf den Namen Ihrer in der Warteschlange befindlichen Ressourcenanfrage.

Nachdem Ihre TPU bereitgestellt wurde, können Sie auch Details zu Ihrer in der Warteschlange befindlichen Ressourcenanfrage aufrufen. Rufen Sie dazu die Seite TPUs auf, suchen Sie nach Ihrer TPU und klicken Sie auf den Namen der entsprechenden in der Warteschlange befindlichen Ressourcenanfrage.

In seltenen Fällen kann es vorkommen, dass die erwartete Ressource den Status FAILED hat, während einige Segmente den Status ACTIVE haben. Löschen Sie in diesem Fall die erstellten Ressourcen und versuchen Sie es nach einigen Minuten noch einmal oder wenden Sie sich an den Google Cloud-Support.

SSH und Abhängigkeiten installieren

Im Hilfeartikel JAX-Code auf TPU-Pod-Slices ausführen wird beschrieben, wie Sie in einem einzelnen Slice eine SSH-Verbindung zu Ihren TPU-VMs herstellen. Wenn Sie über SSH eine Verbindung zu allen TPU-VMs in Ihrer Multislice-Umgebung herstellen und Abhängigkeiten installieren möchten, verwenden Sie den folgenden gcloud-Befehl:

  $ gcloud compute tpus queued-resources ssh ${YOUR_QR_ID} \
        --zone your-zone \
        --node=all \
        --worker=all \
        --command="command-to-run"
        --batch-size=4

Mit diesem gcloud-Befehl wird der angegebene Befehl über SSH an alle Worker und Knoten in QR gesendet. Der Befehl wird in Gruppen von vier zusammengefasst und gleichzeitig gesendet. Die nächste Befehlsgruppe wird gesendet, wenn die Ausführung der aktuellen Gruppe abgeschlossen ist. Wenn bei einem der Befehle ein Fehler auftritt, wird die Verarbeitung beendet und es werden keine weiteren Batches gesendet. Weitere Informationen finden Sie in der API-Referenz für anstehende Ressourcen. Wenn die Anzahl der verwendeten Scheiben das Threadlimit (auch Batching-Limit genannt) Ihres lokalen Computers überschreitet, kommt es zu einem Deadlock. Angenommen, das Batching-Limit auf Ihrem lokalen Computer beträgt 64. Wenn Sie versuchen, ein Trainingsskript auf mehr als 64 Scheiben auszuführen, z. B. 100, werden die Scheiben mit dem SSH-Befehl in Batches aufgeteilt. Das Training wird mit der ersten Gruppe von 64 Scheiben ausgeführt. Danach wird gewartet, bis die Scripts abgeschlossen sind, bevor das Skript auf der verbleibenden Gruppe von 36 Scheiben ausgeführt wird. Die erste Gruppe von 64 Scheiben kann jedoch erst abgeschlossen werden, wenn die verbleibenden 36 Scheiben mit der Ausführung des Scripts beginnen. Dies führt zu einer Deadlock-Situation.

Sie können dieses Szenario vermeiden, indem Sie das Trainingsskript auf jeder VM im Hintergrund ausführen. Fügen Sie dazu dem Skriptbefehl, den Sie mit dem Flag --command angeben, ein kaufmännisches Undzeichen (&) an. Wenn Sie dies tun, wird die Steuerung nach dem Starten des Trainingsscripts für die erste Gruppe von Segmenten sofort an den SSH-Befehl zurückgegeben. Mit dem SSH-Befehl kann dann das Trainingsskript für die verbleibenden 36 Scheiben ausgeführt werden. Sie müssen die stdout- und stderr-Streams entsprechend weiterleiten, wenn Sie die Befehle im Hintergrund ausführen. Um den Parallelismus innerhalb desselben QR-Codes zu erhöhen, können Sie mit dem Parameter --node bestimmte Scheiben auswählen.

Netzwerkeinrichtung

Führen Sie die folgenden Schritte aus, um sicherzustellen, dass TPU-Slices miteinander kommunizieren können. Installieren Sie JAX auf allen Slices. Weitere Informationen finden Sie unter JAX-Code auf TPU-Pod-Slices ausführen. Prüfen Sie, ob len(jax.devices()) der Anzahl der Chips in Ihrer Multislice-Umgebung entspricht. Führen Sie dazu für jeden Ausschnitt Folgendes aus:

  $ python3 -c 'import jax; print(jax.devices())'

Wenn Sie diesen Code auf vier Slices von v4-16 ausführen, gibt es acht Chips pro Slice und vier Slices, sodass jax.devices() insgesamt 32 Chips (Geräte) zurückgeben sollte.

In die Warteschlange gestellte Ressourcen auflisten

gcloud

Mit dem Befehl queued-resources list können Sie den Status Ihrer Ressourcen in der Warteschlange aufrufen:

$ gcloud compute tpus queued-resources list

Die Ausgabe sieht dann ungefähr so aus:

NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central2-b  4           v4-16             ACTIVE
...

Console

1. Rufen Sie in der Google Cloud Console die Seite TPUs auf:

   TPUs aufrufen

2. Klicken Sie auf den Tab In der Warteschlange befindliche Ressourcen.

Job in einer bereitgestellten Umgebung starten

Sie können Arbeitslasten manuell ausführen, indem Sie über SSH eine Verbindung zu allen Hosts in jedem Slice herstellen und den folgenden Befehl auf allen Hosts ausführen.

$ gcloud compute tpus tpu-vm ssh YOUR_QR_ID \
    --zone=your-zone \
    --worker=all \
    --node=all \
    --command="command-to-run"

QR-Codes zurücksetzen

Mit der ResetQueuedResource API können alle VMs in einem ACTIVE-QR-Code zurückgesetzt werden. Durch das Zurücksetzen der VMs wird der Speicher der Maschine gelöscht und die VM auf ihren Ausgangszustand zurückgesetzt. Alle lokal gespeicherten Daten bleiben intakt und das Startskript wird nach dem Zurücksetzen aufgerufen. Die ResetQueuedResource API kann nützlich sein, wenn Sie alle TPUs neu starten möchten. Beispielsweise, wenn das Training nicht weitergeht und das Zurücksetzen aller VMs einfacher ist als das Debuggen.

Das Zurücksetzen aller VMs erfolgt parallel und ein ResetQueuedResource-Vorgang dauert ein bis zwei Minuten. Rufen Sie die API mit dem folgenden Befehl auf:

$ gcloud compute tpus queued-resources reset YOUR_QR_ID

In die Warteschlange gestellte Ressourcen löschen

Wenn Sie Ressourcen am Ende der Trainingseinheit freigeben möchten, löschen Sie die in der Warteschlange befindliche Ressource. Das Löschen dauert zwei bis fünf Minuten. Wenn Sie die gcloud CLI verwenden, können Sie diesen Befehl mit dem optionalen Flag --async im Hintergrund ausführen.

gcloud

$ gcloud compute tpus queued-resources \
    delete YOUR_QR_ID --force (--async)

Console

1. Rufen Sie in der Google Cloud Console die Seite TPUs auf:

   TPUs aufrufen

2. Klicken Sie auf den Tab In der Warteschlange befindliche Ressourcen.

3. Klicken Sie das Kästchen neben der anstehenden Ressourcenanfrage an.

4. Klicken Sie auf Löschen.

Automatische Wiederherstellung nach Fehlern

Bei einer Störung bietet Multislice eine störungsfreie Reparatur des betroffenen Slices und ein anschließendes Zurücksetzen aller Slices. Der betroffene Sliver wird durch einen neuen ersetzt und die übrigen fehlerfreien Sliver werden zurückgesetzt. Wenn keine Kapazität verfügbar ist, um einen Ersatz-Speicherblock zuzuweisen, wird das Training beendet.

Wenn das Training nach einer Unterbrechung automatisch fortgesetzt werden soll, müssen Sie ein Startscript angeben, das nach den zuletzt gespeicherten Checkpoints sucht und sie lädt. Das Startskript wird jedes Mal automatisch ausgeführt, wenn ein Slab neu zugewiesen oder eine VM zurückgesetzt wird. Sie geben ein Startskript in der JSON-Nutzlast an, die Sie an die API zum Erstellen von QR-Codes senden.

Mit dem folgenden Startskript (verwendet in QR-Codes erstellen) können Sie automatisch nach Fehlern wiederherstellen und das Training anhand von Checkpoints fortsetzen, die während des MaxText-Trainings in einem Cloud Storage-Bucket gespeichert wurden:

{
 "tpu": {
   "node_spec": [
     {
      ...
         "metadata": {
               "startup-script": "#! /bin/bash \n pwd \n runuser -l user1 -c 'cd /home/user1/MaxText && python3 MaxText/train.py MaxText/configs/base.yml run_name=run_test_failure_recovery dcn_data_parallelism=4 ici_fsdp_parallelism=8 steps=10000 save_period=10 base_output_directory='gs://user1-us-central2'' EOF"
         }
     ...
     }
   ]
 }
}

Klonen Sie das MaxText-Repository, bevor Sie dies ausprobieren.

Profiling und Fehlerbehebung

Das Profiling ist in Umgebungen mit einer einzelnen und mehreren Scheiben identisch. Weitere Informationen finden Sie unter Profilerstellung für JAX-Programme.

Optimierte Schulungen

Sharding mit Multislice für maximale Leistung

Um in Multi-Slice-Umgebungen die maximale Leistung zu erzielen, müssen Sie überlegen, wie Sie die Shards auf die einzelnen Slices verteilen. Es gibt in der Regel drei Möglichkeiten: Datenparallelität, vollständig shardete Datenparallelität und Pipelineparallelität. Wir empfehlen nicht, Aktivierungsabfragen über die Modelldimensionen zu sharden (manchmal als Tensorparallelität bezeichnet), da dies zu viel Bandbreite zwischen den Slices erfordert. Bei allen diesen Strategien können Sie innerhalb eines Segments dieselbe Sharding-Strategie beibehalten, die in der Vergangenheit für Sie funktioniert hat.

Wir empfehlen, mit reiner Datenparallelität zu beginnen. Die Verwendung von vollständig shardeten Datenparallelismen ist nützlich, um die Arbeitsspeichernutzung zu verringern. Der Nachteil ist, dass die Kommunikation zwischen den Slices über das DCN-Netzwerk erfolgt und die Arbeitslast verlangsamt. Verwenden Sie die Pipeline-Parallelität nur bei Bedarf, je nach Batchgröße (wie unten analysiert).

Wann sollte Datenparallelität verwendet werden?

Die reine Datenparallelität eignet sich gut für Arbeitslasten, die gut funktionieren, deren Leistung Sie aber durch die Skalierung auf mehrere Slices verbessern möchten.

Um eine starke Skalierung über mehrere Slices zu erreichen, muss die Zeit, die für die Ausführung von All-Reduce über das DCN erforderlich ist, kürzer sein als die Zeit, die für die Ausführung eines Rückwärtsdurchlaufs erforderlich ist. Das DCN wird für die Kommunikation zwischen Segmenten verwendet und ist ein begrenzender Faktor für den Arbeitslastdurchsatz.

Jeder TPU v4-Chip erreicht eine Spitzenleistung von 275 * 10¹² FLOPs pro Sekunde.

Es gibt vier Chips pro TPU-Host und jeder Host hat eine maximale Netzwerkbandbreite von 50 Gbit/s.

Die arithmetische Intensität beträgt also 4 × 275 × 10¹² FLOPS ÷ 50 Gbit/s = 22.000 FLOPS ÷ Bit.

Ihr Modell verwendet für jeden Parameter pro Schritt 32 bis 64 Bit DCN-Bandbreite. Wenn Sie zwei Slices verwenden, nutzt Ihr Modell 32 Bit DCN-Bandbreite. Wenn Sie mehr als zwei Slices verwenden, führt der Compiler eine vollständige Zufallsmix-All-Reduce-Operation aus und Sie nutzen für jeden Parameter pro Schritt bis zu 64 Bit DCN-Bandbreite. Die Anzahl der für jeden Parameter erforderlichen FLOPS hängt von Ihrem Modell ab. Insbesondere für transformerbasierte Sprachmodelle beträgt die Anzahl der FLOPS, die für einen Vorwärts- und einen Rückwärtsdurchlauf erforderlich sind, ungefähr 6 * B * P, wobei:

• B ist die Batchgröße in Tokens.
• P ist die Anzahl der Parameter.

Die Anzahl der FLOPS pro Parameter beträgt 6 * B und die Anzahl der FLOPS pro Parameter während des Rückwärtsdurchlaufs 4 * B.

Damit eine effektive Skalierung über mehrere Slices hinweg möglich ist, muss die Betriebsintensität die arithmetische Intensität der TPU-Hardware übersteigen. Um die Betriebsintensität zu berechnen, teilen Sie die Anzahl der FLOPS pro Parameter während des Rückwärtsdurchlaufs durch die Netzwerkbandbreite (in Bits) pro Parameter pro Schritt: Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

Wenn Sie also für ein Transformer-basiertes Sprachmodell zwei Chunks verwenden, gilt Folgendes: Operational intensity = 4 * B / 32

Wenn Sie mehr als zwei Segmente verwenden: Operational intensity = 4 * B/64

Das entspricht einer Mindestbatchgröße von 176.000 bis 352.000 für Transformer-basierte Sprachmodelle. Da das DCN-Netzwerk Pakete kurzzeitig verwerfen kann, sollten Sie einen erheblichen Fehlerspielraum einplanen und die Datenparallelität nur dann implementieren, wenn die Batchgröße pro Pod mindestens 350.000 (zwei Pods) bis 700.000 (viele Pods) beträgt.

Bei anderen Modellarchitekturen müssen Sie die Laufzeit des Rückwärtsdurchlaufs pro Schleife schätzen. Dazu können Sie entweder einen Profiler verwenden oder die FLOPS zählen. Sie können diese Zeit dann mit der erwarteten Ausführungszeit vergleichen, um die Ausführungszeit mit DCN zu reduzieren, und so eine gute Schätzung erhalten, ob der Datenparallelismus für Sie sinnvoll ist.

Wann sollte die vollständig shardierte Datenparallelität (Fully Sharded Data Parallelism, FSDP) verwendet werden?

Bei der vollständigen Sharding-Datenparallelität (Fully Sharded Data Parallelism, FSDP) werden Datenparallelität (Sharding der Daten auf Knoten) und Gewichts-Sharding auf Knoten kombiniert. Für jeden Vorgang im Vorwärts- und Rückwärtsdurchlauf werden die Gewichte gesammelt, damit jeder Stapel die benötigten Gewichte hat. Anstatt die Gradienten mit All-Reduce zu synchronisieren, werden sie beim Erzeugen mit Reduce-Scatter verteilt. So erhält jeder Sliver nur die Gradienten für die Gewichte, für die er verantwortlich ist.

Ähnlich wie bei der Datenparallelität muss bei der FSDP die globale Batchgröße linear mit der Anzahl der Scheiben skaliert werden. Mit FSDP wird der Speicherdruck verringert, wenn Sie die Anzahl der Scheiben erhöhen. Das liegt daran, dass die Anzahl der Gewichte und der Optimizer-Status pro Slither verringert wird. Dies geht jedoch zu Lasten des Netzwerkverkehrs und erhöht die Wahrscheinlichkeit von Blockierungen aufgrund eines verzögerten Kollektivs.

In der Praxis eignet sich die FSDP-Methode über mehrere Slices hinweg am besten, wenn Sie den Batch pro Slice erhöhen, mehr Aktivierungen speichern, um die Neumaterialisierung während des Rückwärtsdurchlaufs zu minimieren, oder die Anzahl der Parameter in Ihrem neuronalen Netzwerk erhöhen.

Die All-Gather- und All-Reduce-Vorgänge in FSDP funktionieren ähnlich wie in DP. Sie können also auf die gleiche Weise wie im vorherigen Abschnitt feststellen, ob Ihre FSDP-Arbeitslast durch die DCN-Leistung begrenzt ist.

Wann sollte Pipelineparallelität verwendet werden?

Der Pipeline-Parallelismus wird relevant, wenn Sie mit anderen Parallelisierungsstrategien eine hohe Leistung erzielen, für die eine globale Batchgröße erforderlich ist, die über Ihrer bevorzugten maximalen Batchgröße liegt. Durch die Pipelineparallelität können die Segmente einer Pipeline einen Batch „teilen“. Die Pipeline-Parallelität hat jedoch zwei erhebliche Nachteile:

1. Es kommt zu einer „Pipeline-Blase“, bei der Chips inaktiv sind, weil sie auf Daten warten.
2. Es erfordert Mikro-Batching, wodurch die effektive Batch-Größe, die arithmetische Intensität und letztendlich die FLOP-Nutzung des Modells verringert werden.

Der Pipeline-Parallelismus sollte nur verwendet werden, wenn die anderen Parallelitätsstrategien eine zu große globale Batchgröße erfordern. Bevor Sie die Pipelineparallelität ausprobieren, sollten Sie empirisch testen, ob sich die Konvergenz pro Stichprobe bei der Batchgröße verlangsamt, die für eine hohe FSDP-Leistung erforderlich ist. Bei der FSDP wird in der Regel eine höhere FLOP-Nutzung des Modells erreicht. Wenn sich die Konvergenz pro Stichprobe jedoch mit zunehmender Batchgröße verlangsamt, ist der Pipeline-Parallelismus möglicherweise die bessere Wahl. Die meisten Arbeitslasten können ausreichend große Batchgrößen tolerieren, um nicht von der Pipelineparallelität zu profitieren. Bei Ihrer Arbeitslast kann das jedoch anders sein.

Wenn Pipelineparallelität erforderlich ist, empfehlen wir, sie mit Datenparallelität oder FSDP zu kombinieren. So können Sie die Pipelinetiefe minimieren und gleichzeitig die Batchgröße pro Pipeline erhöhen, bis die DCN-Latenz weniger Einfluss auf den Durchsatz hat. Wenn Sie beispielsweise N Slices haben, sollten Sie Pipelines mit einer Tiefe von 2 und N/2 Repliken der Datenparallelität, dann Pipelines mit einer Tiefe von 4 und N/4 Repliken der Datenparallelität usw. in Betracht ziehen, bis der Batch pro Pipeline groß genug ist, dass die DCN-Kollektive hinter der Arithmetik im Rückwärtsdurchlauf verborgen werden können. Dadurch wird die durch den Pipelineparallelismus verursachte Verlangsamung minimiert und Sie können über das globale Limit für die Batchgröße hinaus skalieren.

Best Practices für Mehrfachaufnahmen

Laden der Daten

Während des Trainings laden wir wiederholt Batches aus einem Dataset, um sie in das Modell einzugeben. Ein effizienter, asynchroner Datenlader, der den Batch auf mehrere Hosts verteilt, ist wichtig, damit die TPUs nicht zu wenig Arbeit haben. Beim aktuellen Datenloader in MaxText lädt jeder Host eine gleiche Teilmenge der Beispiele. Diese Lösung ist für Text geeignet, erfordert aber eine Neuaufteilung innerhalb des Modells. Außerdem bietet MaxText noch keine deterministischen Snapshots, mit denen der Dateniterator dieselben Daten vor und nach der Voraktivierung laden könnte.

Prüfpunktausführung

Die Checkpoint-Bibliothek Orbax bietet Primitive zum Erstellen von Checkpoints für JAX-PyTrees im lokalen Speicher oder in Google Cloud Storage. Wir stellen eine Referenzintegration mit synchronem Checkpointing in MaxText in checkpointing.py bereit.

Unterstützte Konfigurationen

Formen

Alle Scheiben müssen dieselbe Form haben (z. B. dieselbe AcceleratorType). Heterogene Scheibenformen werden nicht unterstützt.

Orchestrierung

Die Orchestration wird mit GKE unterstützt. Weitere Informationen finden Sie unter TPUs in GKE.

Frameworks

Multislice unterstützt nur JAX- und PyTorch-Arbeitslasten.

Parallelität

Wir empfehlen Nutzern, Multislice mit Datenparallelität zu testen. Weitere Informationen zur Implementierung von Pipeline-Parallelität mit Multislice erhalten Sie von Ihrem Google Cloud-Kundenbetreuer.

Support und Feedback

Wir freuen uns über jedes Feedback. Wenn Sie Feedback geben oder Support anfordern möchten, verwenden Sie bitte das Cloud TPU-Support- oder Feedbackformular.