Cloud TPU-Multislice-Übersicht [öffentliche Vorschau]

Vorschau

Dieses Produkt oder diese Funktion unterliegt den „Nutzungsbedingungen für Pre-GA-Angebote“ im Abschnitt „Allgemeine Nutzungsbedingungen“ der dienstspezifischen Nutzungsbedingungen. Pre-GA-Produkte und -Funktionen sind unverändert verfügbar und werden möglicherweise nur eingeschränkt unterstützt. Weitere Informationen finden Sie in den Beschreibungen der Startphase.

Cloud TPU Multislice ist eine Full-Stack-Technologie zur Leistungsskalierung, mit der ein Trainingsjob mehrere TPU-Slices innerhalb eines einzelnen Pods oder auf Slices in mehreren Pods mit einfacher Datenparallelität verwenden kann. Mit TPU v4-Chips können Trainingsjobs mehr als 4.096 Chips in einer einzigen Ausführung verwenden. Bei Trainingsjobs, die weniger als 4.096 Chips erfordern, kann ein einzelnes Segment die beste Leistung bieten. Mehrere kleinere Segmente sind jedoch leichter verfügbar. Dies ermöglicht eine schnellere Startzeit, wenn Multislice mit kleineren Segmenten verwendet wird.

Bei der Bereitstellung in Multislice-Konfigurationen kommunizieren TPU-Chips in jedem Segment über Inter-Chip-Interconnect (ICI). TPU-Chips in verschiedenen Slices kommunizieren durch die Übertragung von Daten an CPUs (Hosts), die wiederum die Daten über das Rechenzentrumsnetzwerk (DCN) übertragen.

Entwickler schreiben keinen Code für die Kommunikation zwischen DCNs. Der XLA-Compiler generiert diesen Code für Sie und überschneidet die Kommunikation mit der Berechnung für maximale Leistung.

Konzepte

AcceleratorType

Die Form jedes TPU-Slices, das ein Multislice umfasst. Jedes Segment in einer Multiplikationsanfrage hat den gleichen AcceleratorType. AcceleratorType besteht aus einem TPU-Typ (v2, v3, v4, v5e und höher), gefolgt von der Anzahl der TensorCores. Beispiel: v4-128 gibt eine TPU v4 mit 128 TensorCores an.

Automatische Reparatur

Wenn bei einem Segment ein Wartungsereignis, ein vorzeitiges Beenden oder ein Hardwarefehler auftritt, erstellt Cloud TPU ein neues Segment. In den seltenen Fällen, in denen nicht genügend Ressourcen vorhanden sind, um ein neues Segment zu erstellen, wird das Erstellen erst dann abgeschlossen, wenn die Hardware verfügbar ist. Anschließend werden alle anderen Segmente in der Multislice-Umgebung neu gestartet, damit das Training fortgesetzt werden kann. Mit einem ordnungsgemäß konfigurierten Startskript kann das Trainingsskript automatisch neu gestartet werden, ohne dass der Nutzer eingreifen muss. Es wird ab dem letzten Prüfpunkt geladen und fortgesetzt.

Dataset

Die Daten, die von einem Modell für Training oder 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.

Planung der Gang

Wenn alle TPU-Slices zusammen bereitgestellt werden, wird entweder garantiert, dass alle oder keine Slices erfolgreich bereitgestellt werden.

Host

Ein Host ist ein physischer Computer, auf dem VMs ausgeführt werden. Ein Host kann maximal vier VMs gleichzeitig ausführen. Jede VM hat eine dedizierte TPU.

Inferenz

Ein vortrainiertes Modell für maschinelles Lernen auf einem Host laden und Vorhersagen zu Daten treffen.

Interchip Interconnect (ICI)

Interne Links mit hoher Geschwindigkeit und niedriger Latenz, die TPUs in einem TPU-Pod verbinden.

Multislice

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

Knoten

Im Kontext Multislice bezieht sich der Knoten auf einen einzelnen TPU-Slice. Jedes TPU-Segment in einem Multislice erhält eine Knoten-ID.

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.

Warteschlange (QR)

Eine Darstellung von TPU-Ressourcen, die zum Einstellen und Verwalten einer Anfrage für eine Single-Slice- oder Multislice-TPU-Umgebung verwendet werden.

Startskript

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

TPU-Slice

Ein logischer Unterabschnitt eines TPU-Pod mit TPU-Chips. Alle Chips in einem Segment kommunizieren über das ICI-Netzwerk miteinander.

TPU-VM

Eine virtuelle Maschine, auf der Linux ausgeführt wird und Zugriff auf die zugrunde liegenden TPUs hat. Bei v4-TPUs hat jede TPU-VM direkten Zugriff auf vier Chips. Manchmal rufen wir eine TPU-VM als Worker auf.

Google Tensor

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

Tensor Processing Unit (TPU)

Der intern entwickelte ML-Beschleunigungschip von Google. Sie wurden entwickelt, um die schnellste und effizienteste Rechenleistung für wichtige ML-Aufgaben wie die Matrixmultiplikation zu bieten.

Arten von Cloud TPU-Kapazitäten

TPUs können aus drei Kapazitätstypen erstellt werden (siehe Nutzungsoptionen unter Wie TPU-Preise funktionieren) :

• Reservierung: Ausrichtung auf reserviertes Kontingent. Um ein reserviertes Kontingent zu verwenden, müssen Sie eine Reservierungsvereinbarung mit Google haben. Verwenden Sie beim Erstellen Ihrer Ressourcen das Flag --reserved.
• Auf Abruf: Ausrichtung auf das Kontingent auf Abruf. Ihre Ressourcen werden möglicherweise vorzeitig beendet, um Platz für Anfragen für einen Job mit höherer Priorität zu schaffen. Verwenden Sie beim Erstellen Ihrer Ressourcen das Flag --best-effort.
• On-Demand: Ausrichtung auf On-Demand-Kontingent. Es ist keine Reservierung erforderlich und wird nicht vorzeitig beendet. Die TPU-Anfrage wird in eine On-Demand-Kontingentwarteschlange von Cloud TPU gestellt. Standardmäßig ausgewählt, keine Flags erforderlich.

Hinweis: Cloud TPU begrenzt die Segmentgröße für alle Beschleunigertypen. Informationen zu Slice-Größengrößen für On-Demand-Anfragen finden Sie unter Standard-Slice-Größen für On-Demand-Ressourcen in der Warteschlange.

Jetzt starten

Wenn Sie noch keine TPUs verwendet haben, installieren Sie zuerst die Komponenten gcloud und aktivieren Sie die Cloud TPU API. Konfigurieren Sie dann den Befehl gcloud, um Ihr Projekt festzulegen. Folgen Sie der Anleitung unter Projekt einrichten und beenden Sie sie, wenn Sie den Abschnitt TPUs verwalten erreicht haben.

Zum Verwenden von Multislice müssen Ihre TPU-Ressourcen als Ressourcen in der Warteschlange verwaltet werden. Wenn Sie bereits TPU v4-Nutzer sind und eine Reservierung haben, müssen Sie Ihre Reservierung möglicherweise zu einem neuen Reservierungssystem migrieren. Weitere Informationen erhalten Sie vom Cloud TPU-Team. In GA kann die Anfrage zum Erstellen von Ressourcen in der Warteschlange auf alle Reservierungen zugreifen und Nutzer müssen keine Kapazität übertragen, bevor sie Multislice-Ressourcen oder Ressourcen in der Warteschlange verwenden können.

Einführungsbeispiel

In dieser Anleitung wird Code im MaxText-GitHub-Repository verwendet. shardings.py (im MaxText-Repository) soll das Experimentieren mit verschiedenen Optionen für die Parallelität vereinfachen. Dieser Dienst ist der beste Ausgangspunkt, wenn Sie mehr über Datenparallelität, Fully Sharded Data Parallelism (FSDP) und/oder Tensor-Parallelität erfahren möchten. Mit diesem Code können Sie den Code sowohl auf einzelne Segmente als auch auf mehrere Segmente erweitern.

Dies ist nicht spezifisch für Multiplikationen und sollte an anderer Stelle leben.

Einteiliges Fragmentieren mit ICI-Parallelität

Die ICI-Fragmentierung entspricht der Fragmentierung in jedem Segment. shardings.py bietet drei ICI-Parallelitätseingaben:

• 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 Parallelitätsmethode.

Diese Eingaben müssen eingeschränkt werden, sodass ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism der Anzahl der Chips im Segment entspricht.

Die folgende Tabelle zeigt Beispieleingaben für ICI-Parallelität für die vier Chips, die in v4-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
2-Wege-FSDP + 2-Wege-Tensor-Parallelität	1	2	2

Beachten Sie, dass ici_data_parallelism in den meisten Fällen 1 bleiben sollte, da das ICI-Netzwerk schnell genug ist, um FSDP fast immer einer Datenparallelität vorzuziehen.

In diesem Beispiel wird davon ausgegangen, dass Sie mit dem Ausführen von Code auf einem einzelnen TPU-Slice vertraut sind, z. B. unter Berechnung auf einer Cloud TPU-VM mit JAX ausführen. Dieses Beispiel wurde aus dieser Anleitung angepasst, um shardings.py auf einem einzelnen Segment auszuführen.

1. Umgebung einrichten:

   $ QUEUED_RESOURCE_ID=user-defined-queued-resource-id
   PROJECT=your-gcp-project-name
   ZONE=us-central2-b
   RUNTIME_VERSION=tpu-ubuntu2204-base
   
   gcloud auth login
   gcloud config set project ${PROJECT}
   gcloud config set compute/zone ${ZONE}
   

2. Erstellen Sie SSH-Schlüssel für gcloud. Wir empfehlen, ein leeres Passwort anzugeben. Drücken Sie nach der Ausführung des folgenden Befehls zweimal „Eingabetaste“. Wenn Sie gefragt werden, ob die Datei google_compute_engine bereits vorhanden ist, ersetzen Sie die vorhandene Version.

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

3. Stellen Sie Ihre TPUs mit dem folgenden Befehl bereit:

   $ gcloud alpha compute tpus queued-resources create $QUEUED_RESOURCE_ID \
   --accelerator-type v4-8 \
   --runtime-version tpu-ubuntu2204-base \
   --node-id $QUEUED_RESOURCE_ID \
   [--reserved |--best-effort]
   

   Hinweis:Derzeit unterstützt die gcloud-CLI nicht alle Optionen zum Erstellen von Ressourcen in der Warteschlange, z. B. Tags. Weitere Informationen finden Sie unter Ressourcen in der Warteschlange erstellen.

4. Warten Sie, bis sich der QR-Code im Status ACTIVE befindet, was bedeutet, dass die Worker-Knoten den Status READY haben. Nachdem die QR-Bereitstellung gestartet wurde, kann der Vorgang je nach Größe des QR-Codes ein bis fünf Minuten dauern. Sie können den Status mit dem folgenden Befehl prüfen:

   $ gcloud alpha compute tpus queued-resources \
   list --filter=$QUEUED_RESOURCE_ID
   

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

   $ gcloud compute tpus tpu-vm ssh $QUEUED_RESOURCE_ID
   

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

7. Führen Sie im MaxText-Repository-Verzeichnis das Einrichtungsskript aus, um JAX und andere Abhängigkeiten auf Ihrem TPU-Slice zu installieren.

   $ bash setup.sh # These installations will take
   a few minutes.
   

8. Führen Sie den folgenden Befehl aus, um shardings.py im TPU-Slice auszuführen.

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

   Die Ergebnisse werden in den Logs angezeigt. Ihre TPUs sollten etwa 260 TFLOP/s erreichen oder eine beeindruckende FLOP-Auslastung von über 90%haben. In diesem Fall haben wir den maximalen Batch ausgewählt, der in den Arbeitsspeicher der hohen Bandbreite (HBM) passt.

9. Sie können andere Fragmentierungsstrategien für ICI nutzen, z. B. die folgende Kombination:

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

10. Löschen Sie den QR- und TPU-Slice, wenn Sie fertig sind. Sie sollten diese Bereinigungsschritte in der Umgebung ausführen, in der Sie das Segment eingerichtet haben. Führen Sie zuerst exit aus, um die SSH-Sitzung zu beenden. Der Löschvorgang dauert zwei bis fünf Minuten und kann im Hintergrund mit dem optionalen Flag --async ausgeführt werden.

    Wichtig :QR-Codes müssen nach Gebrauch gelöscht werden. Andernfalls verbrauchen sie weiterhin Ihr QR-Kontingent, das vom TPU-Kontingent getrennt erfasst wird.

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

Multislice-Fragmentierung mit DCN-Parallelität

Das Skript shardings.py nimmt drei Parameter, die die DCN-Parallelität angeben, entsprechend der Anzahl der Shards jeder Art von Datenparallelität:

• dcn_data_parallelism
• dcn_fsdp_parallelismus
• dcn_tensor_parallelism

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

Verwenden Sie als Beispiel für zwei Segmente --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 gesetzt werden, weil DCN für solche Fragmentierung nicht geeignet ist. Bei typischen LLM-Arbeitslasten auf v4-Chips sollte dcn_fsdp_parallelism auf 1 gesetzt werden. Daher sollte für dcn_data_parallelism die Anzahl der Segmente festgelegt werden, dies ist jedoch abhängig von der Anwendung.

Wenn Sie die Anzahl der Segmente erhöhen (unter Berücksichtigung der Segmentgröße und des Batchwerts pro Segment), erhöhen Sie gleichzeitig die Parallelität der Daten. Eine Anleitung finden Sie auch in der Readme für MaxText.

shardings.py wird in einer Multislice-Umgebung ausgeführt

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 den Schritten aus Erste Schritte: Mehrere Tests für mehrere Segmente aus dem MaxText-Repository. Hier führen wir jedoch shardings.py anstelle des komplexeren LLM in train.py aus.

Das multihost_runner.py-Tool ist für schnelle Tests optimiert, wobei dieselben TPUs wiederholt wiederverwendet werden. Da das Skript multihost_runner.py von langlebigen SSH-Verbindungen abhängt, wird es für Jobs mit langer Ausführungszeit nicht empfohlen. Wenn Sie einen längeren Job ausführen möchten (z. B. Stunden oder Tage), empfehlen wir die Verwendung von multihost_job.p.

In dieser Anleitung wird der Begriff Runner verwendet, um den Computer anzugeben, auf dem das multihost_runner.py-Skript ausgeführt wird. Mit dem Begriff Worker werden die TPU-VMs angegeben, aus denen die Segmente bestehen. Sie können multihost_runner.py auf einer lokalen Maschine oder einer beliebigen Compute Engine-VM im selben Projekt wie Ihre Segmente 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 werden shardings.py über zwei v4-16-Slices ausgeführt, insgesamt also 4 VMs und 16 TPU-Chips. Sie können das Beispiel ganz einfach ändern, damit es auf mehr TPUs ausgeführt wird.

Umgebung einrichten

1. Klonen Sie MaxText auf Ihrem Runner-Computer.
2. Wechseln Sie zum Repository-Verzeichnis.
3. Umgebung einrichten:

   $ export QUEUED_RESOURCE_ID=user-defined-queued-resource-id
   $ export PROJECT=your-gcp-project-name
   $ export ZONE=us-central2-b
   $ export RUNTIME_VERSION=tpu-ubuntu2204-base
   $ gcloud config set project ${PROJECT}
   $ gcloud config set compute/zone ${ZONE}
   

4. Erstellen Sie SSH-Schlüssel für gcloud. Wir empfehlen, ein leeres Passwort anzugeben. Drücken Sie nach der Ausführung des folgenden Befehls zweimal die Eingabetaste. Wenn Sie gefragt werden, ob die Datei google_compute_engine bereits vorhanden ist, behalten Sie die vorhandene Version bei.

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

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

   $ export SLICE_COUNT=2
   

6. Erstellen Sie mit queued-resources create eine Multislice-Umgebung.

   Hinweis: Achten Sie darauf, dass Sie das entsprechende Kontingent haben, bevor Sie --reserved, --best_effort oder das standardmäßige On-Demand-Kontingent auswählen. Informationen zu Kontingenttypen finden Sie unter Kontingentrichtlinie.

   Der folgende Befehl zeigt, wie Sie eine Multislice TPU aus v4 erstellen. TPU v5e wird auch in der öffentlichen Multislice-Vorschau unterstützt. Geben Sie einen v5e accelerator-type (z. B. v5e-16) und einen v5e runtime-version (v2-alpha-tpuv5-lite) an, um v5e zu verwenden.

   $ gcloud alpha compute tpus queued-resources \
   create $QUEUED_RESOURCE_ID \
   --accelerator-type=v4-16 \
   --runtime-version=tpu-ubuntu2204-base \
   --node-count=$SLICE_COUNT \
   --node-prefix=$QUEUED_RESOURCE_ID \
   [--reserved|--best-effort]
   

   Beschreibung der Befehls-Flags

   QUEUED_RESOURCE_ID

   Ein benutzerdefinierter String, der die Anfrage in der Warteschlange identifiziert.

   accelerator-type

   Der Typ der zu erstellenden Cloud TPU.

   runtime-version

   Die [Cloud TPU-Softwareversion](/tpu/docs/supported-tpu-configurations#tpu_software_versions).

   node-count

   Die Anzahl der zu erstellenden Segmente.

   node-prefix

   Das Präfix, mit dem die Namen für die einzelnen Segmente generiert werden. An das Präfix wird für jedes Segment eine Zahl angehängt. Wenn Sie beispielsweise node-prefix auf mySlice festlegen, erhalten die Segmente den Namen mySlice-0, mySlice-1 usw.

   reserved

   Beim Erstellen der Segmente können Sie ein reserviertes Kontingent verwenden.

   best-effort

   Beim Erstellen der Segmente Best-Effort-Kontingent verwenden [Standard].

7. Warten Sie, bis sich die Warteschlange der Ressourcen in der Warteschlange (ACTIVE) befindet, was bedeutet, dass sich die Worker-Knoten im Status READY befinden. Nachdem die QR-Bereitstellung gestartet wurde, kann der Vorgang je nach Größe des QR-Codes ein bis fünf Minuten dauern. Sie können den Status mit dem folgenden Befehl prüfen:

   $ gcloud alpha compute tpus queued-resources list \
   --filter=$QUEUED_RESOURCE_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
   ...
   

   Kontaktieren Sie uns, wenn der QR-Status länger als 15 Minuten im Status ACCEPTED (WAITING_FOR_RESOURCES) oder PROVISIONING ist.

8. Installieren Sie die Abhängigkeiten:

   $ pip3 install absl-py
   $ python3 multihost_runner.py \
     --TPU_PREFIX=$QUEUED_RESOURCE_ID \
     --COMMAND="bash setup.sh"
   

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

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

   Die Protokolldateien enthalten etwa 230 TFLOPs/s.

10. Reinigen Sie die TPUs und den QR-Code, wenn Sie fertig sind. Das Löschen dauert zwei bis fünf Minuten und kann im Hintergrund mit dem optionalen Flag --async ausgeführt werden.

MaxText: Leistungsstarker, skalierbarer Multislice LLM

Mit train.py haben wir eine leistungsstarke, skalierbare, lesbare LLM-Implementierung bereitgestellt, die sofort einsetzbar ist. Wir empfehlen Ihnen, es selbst auszuführen und den Code in Ihren eigenen Implementierungen zu verwenden. MaxText enthält eine eigene Readme-Datei und Leistungsergebnisse.

Die Ergebnisse von MaxText sind enthalten und wurden aus der MaxText-Readme-Datei kopiert. MaxText enthält auch alles, was die Ergebnisse reproduzieren kann. Hier gehen wir von FLOPS bei 6 * BATCH * PARAMETER aus.

22 Milliarden Parametermodell

Hardware	TFLOP/Sek./Chip	Modell FLOP Auslastung
1x v4-128	165	60%
2x v4–128	158	57,4%
4X V4-128	157	57,1%

52 Milliarden Parametermodell

Hardware	TFLOP/Sek./Chip	Modell: FLOP-Auslastung (MFU)
1x v4-384	154	56,0%
2x v4–384	162	58,9%

Die Verbesserungen von 1x v4-384 auf 2x v4-384 sind auf eine Compiler-Quirke zurückzuführen. Möglicherweise sehen Sie sie unterschiedlich

Arbeitslast auf Multislice skalieren

Änderungen am Code

Bevor Sie das Modell in einer Multislice-Umgebung ausführen, müssen Sie den folgenden Code ändern:

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

• Legen Sie einige Umgebungsvariablen fest:

  os.environ["JAX_USE_PJRT_C_API_ON_TPU"] = "1"
  os.environ["LIBTPU_INIT_ARGS"] = "--xla_tpu_enable_megascale_barrier=true"
  

  Hinweis: LIBTPU_INIT_ARGS ist eine durch Leerzeichen getrennte Liste, z. B. LIBTPU_INIT_ARGS="--flag1=val1 --flag2=var2".

Dies sollte die einzige erforderliche Codeänderung beim Wechsel zu Multislice sein. Um eine hohe Leistung zu erzielen, muss DCN den parallelen, vollständig fragmentierten Daten der parallelen Achsen oder parallelen Pipelineachsen zugeordnet werden. Weitere Informationen zu Leistungsaspekten und Fragmentierungsstrategien finden Sie unter Fragmentierung mit Multislice für maximale Leistung.

Um zu validieren, dass Ihr Code auf alle Geräte zugreifen kann, können Sie sicher sein, dass len(jax.devices()) der Anzahl der Chips in Ihrer Multislice-Umgebung entspricht. Wenn Sie beispielsweise vier Segmente von v4-16 verwenden, haben Sie acht Chips pro Segment * 4 Segmente. len(jax.devices()) sollte also 32 zurückgeben.

Slice-Größen für Multislice-Umgebungen auswählen

Fügen Sie neue Segmente mit der gleichen Größe wie das bestehende Segment hinzu, um eine lineare Geschwindigkeit zu erhalten. Wenn Sie beispielsweise derzeit ein Segment vom Typ v4-512 verwenden, erreicht Multislice die Leistung etwa doppelt so stark. Dazu wird ein zweites Segment des Typs v4-512 hinzugefügt und die globale Batchgröße verdoppelt. Weitere Informationen finden Sie unter Fragmentierung mit Multislice für maximale Leistung.

Job auf mehreren Slices ausführen

Es gibt drei verschiedene Ansätze zum Ausführen Ihrer benutzerdefinierten Arbeitslast in einer Multislice-Umgebung:

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

Skript für die Ausführung des Tests

Für schnelle Tests stellen Sie multihost_runner.py (in MaxText) bereit, ein einfaches Skript zum Verteilen von Code an eine Multislice-Umgebung, zum Ausführen des Befehls auf jedem Host, zum Kopieren der Logs und zum Erfassen des Fehlerstatus der einzelnen Befehle. Das Skript multihost_runner.py ist in der Readme-Datei für MaxText ausführlicher beschrieben.

Da multihost_runner.py dauerhafte SSH-Verbindungen aufrechterhält, ist sie nur für Tests mit relativ geringer Größe und relativ kurzer Laufzeit geeignet. Passen Sie die Schritte in der Anleitung zu Multihost_Runner.py an Ihre Arbeitslast- und Hardwarekonfiguration an.

Skript für die Produktionsausführung

Bei Produktionsjobs, die Ausfallsicherheit bei Hardwarefehlern und anderen vorzeitigen Ausfällen benötigen, empfiehlt sich die direkte Einbindung in die Create Queued Resource API. Als funktionierendes Beispiel stellen wir multihost_job.py (in MaxText gehostet) bereit, das den Aufruf der Created Queued Resource API mit dem entsprechenden Startskript auslöst, um das Training auszuführen und mit dem vorzeitigen Beenden fortzufahren. Das Skript multihost_job.py ist in der Readme-Datei für MaxText ausführlicher beschrieben.

Da multihost_job.py für jede Ausführung Ressourcen bereitstellen muss, bietet dies nicht so einen Iterationszyklus wie multihost_runner.py.

Manueller Ansatz

Wir empfehlen, dass Sie multihost_Runner.py oder multihost_job.py verwenden oder anpassen, um Ihre benutzerdefinierte Arbeitslast in der gewünschten Multislice-Konfiguration auszuführen. Wenn Sie Ihre Umgebung jedoch mit Warteschlangenressourcen direkt bereitstellen und verwalten möchten, lesen Sie Multislice-Umgebung verwalten.

Multislice-Umgebung verwalten

Wenn Sie Ressourcen in der Warteschlange manuell bereitstellen und verwalten möchten, ohne die im MaxText-Repository enthaltenen Tools zu verwenden, lesen Sie die folgenden Abschnitte.

Ressourcen in Warteschlange

Legen Sie die folgenden Umgebungsvariablen vor der Bereitstellung der Kapazität fest:

  $ export QUEUED_RESOURCE_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

Eingabe	Beschreibung
QUEUED_RESOURCE_ID	Die vom Nutzer zugewiesene ID der Warteschlange.
PROJEKT	Google Cloud-Projektname
ZONE	us-central2-b
NETZWERK_NAME	Name der VPC-Netzwerke.
UNTERNETZWERK_NAME	Name des Subnetzes in VPC-Netzwerken
LAUFZEIT	tpu-ubuntu2204-Basis
ACCELERATOR_TYPE	V4-16
EXAMPLE_TAG_1, EXAMPLE_TAG_2...	Tags zur Identifizierung gültiger Quellen oder Ziele für Netzwerkfirewalls
SLICE_COUNT	Anzahl der Segmente. Beschränkt auf maximal 32 Segmente, kann aber während der öffentlichen Vorschau erhöht werden.
STARTUP_SKRIPT	Wird der Erstellungsanfrage hinzugefügt, kann ein Startskript ausgeführt werden, wenn ein TPU-Slice bereitgestellt/neu gestartet wird und das TPU-Slice repariert oder zurückgesetzt wird.

Mit gcloud eine Ressourcenanforderung in die Warteschlange stellen

$ gcloud alpha compute tpus queued-resources create ${QUEUED_RESOURCE_ID} \
  --project ${PROJECT} \
  --zone ${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|--best-effort]
  

Sorgen Sie dafür, dass Sie das entsprechende Kontingent haben, bevor Sie --reserved, --best_effort oder das standardmäßige On-Demand-Kontingent auswählen. Informationen zu Kontingenttypen finden Sie unter Kontingentrichtlinie.

Mit curl eine Ressourcenanforderung in die Warteschlange stellen

Erstellen Sie eine Datei mit dem Namen queuedresourcereq.json und kopieren Sie den folgenden JSON-Code in diese Datei.

{
  "guaranteed": { "reserved": true },
  "tpu": {
    "node_spec": [
      {
        "parent": "projects/<PROJECT>/locations/<ZONE>",
        "node": {
          "accelerator_type": "<ACCELERATOR_TYPE>",
          "runtime_version": "<RUNTIME_VERSION>",
          "network_config": {
            "network": "<NETWORK_NAME>",
            "subnetwork": "<SUBNETWORK_NAME>",
            "enable_external_ips": true
          },
     "tags" : ["<EXAMPLE_TAG_1>","<EXAMPLE_TAG_2>"]
          "metadata": {
            "startup-script": "<STARTUP_SCRIPT>"
          }
        },
        "multi_node_params": {
          "node_count": <SLICE_COUNT>,
          "node_id_prefix": "<QUEUED_RESOURCE_ID>"
        }
      }
    ]
  }
}

Ersetzen Sie die Elemente in "< >" durch die entsprechenden Werte. Weitere Informationen finden Sie in der Dokumentation zur REST Queued Resource API.

Hinweis: Obwohl node_spec eine Liste ist, wird derzeit nur ein node_spec-Eintrag unterstützt.

Um die erwerbbare Kapazität zu nutzen, ersetzen Sie:

"guaranteed": { "reserved": true }

mit

"best_effort": {}

Alternativ können Sie die Zeile entfernen, um die standardmäßige On-Demand-Kapazität zu verwenden.

Senden Sie die Anfrage zum Erstellen eines QR-Codes mit der JSON-Nutzlast:

  $ 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/${PROJECT}/locations/${ZONE}/queuedResources\?queued_resource_id\=${QUEUED_RESOURCE_ID}

Sie sollten eine Ausgabe wie die folgende mit dem Vorgangsnamen für die QR-Erstellungsanfrage sehen:

{
  "name": "projects/123456789012/locations/us-central2-b/operations/operation-0000000000000-0000000000000-00000000-00000000",
  "metadata": {
    "@type": "type.googleapis.com/google.cloud.common.OperationMetadata",
    "createTime": "2023-03-09T07:35:11.142588587Z",
    "target": "projects/123456789012/locations/us-central2-b/queuedResources/your-queued-resource-id",
    "verb": "create",
    "cancelRequested": false,
    "apiVersion": "v2alpha1"
  },
  "done": false
}

Sie können den Vorgang name verwenden, um den Status des QR-Erstellen-Befehls abzurufen.

Vorgang abrufen

Verwenden Sie Folgendes, um die Details des Vorgangs zum Erstellen eines QR-Codes zu prüfen:

  $ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2alpha1/projects/$PROJECT/locations/$ZONE/operations/operation-00000000000000-0000000000000-00000000-00000000

Die Antwort dieses Befehls enthält den Status des Vorgangs:

{
  "name": "projects/project/locations/zone/operations/operation-00000000000000-0000000000000-00000000-00000000",
  "metadata": {...},
  "done": true,
  "response": {
    "@type": "type.googleapis.com/google.cloud.tpu.v2alpha1.QueuedResource",
    ...
    "state": {
      "state": "WAITING_FOR_RESOURCES"
    }
  }
}

Wenn der QR-Code ("done = true") erstellt wurde, ist der Status im Feld response entweder ACCEPTED (WAITING_FOR_RESOURCES) oder FAILED. Wenn sich der QR-Code im Status ACCEPTED (WAITING_FOR_RESOURCES) befindet, wurde der QR-Code in die Warteschlange gestellt und beginnt mit der Bereitstellung, sobald genügend Ressourcen vorhanden sind. Wenn der QR-Code den Status FAILED hat, wird der Grund für den Fehler in der Ausgabe angezeigt. Weitere Informationen zu anderen möglichen Status finden Sie im Nutzerhandbuch für Ressourcen in der Warteschlange.

Verwenden Sie nach Abschluss des Vorgangs die Ressourcen in der Warteschlange, um die Phasen der QR-Datei zu überwachen.

In seltenen Fällen sehen Sie möglicherweise den QR-Status FAILED, während einige Segmente ACTIVE sind. Löschen Sie in diesem Fall die erstellten Ressourcen und versuchen Sie es in einigen Minuten noch einmal oder wenden Sie sich an das Cloud TPU-Team, um das Problem zu beheben.

SSH-Verbindung installieren und Abhängigkeiten installieren

Unter JAX-Code auf TPU-Pod-Slices ausführen wird beschrieben, wie Sie eine SSH-Verbindung zu Ihren TPU-VMs in einem einzelnen Segment herstellen. Verwenden Sie den folgenden gcloud-Befehl, um eine Verbindung zu allen TPU-VMs in Ihrer Multislice-Umgebung über SSH herzustellen und Abhängigkeiten zu installieren:

gcloud compute tpus queued-resources ssh ${QUEUED_RESOURCE_ID} \
  --zone ${ZONE} \
  --node=all \
  --worker=all \
  --command="some-command"
  --batch-size=4

Mit dem Befehl gcloud wird der angegebene Befehl per SSH an alle Worker und Knoten in QR gesendet. Der Befehl wird in Vierergruppen zusammengefasst und gleichzeitig gesendet. Der nächste Befehlssatz wird gesendet, wenn der aktuelle Batch die Ausführung abgeschlossen hat. Wenn bei einem der Befehle ein Fehler auftritt, werden die Verarbeitungen beendet und es werden keine weiteren Batches gesendet. Weitere Informationen finden Sie in der API-Referenz zur Warteschlange mit Ressourcen.

Netzwerkeinrichtung

Führen Sie die folgenden Schritte aus, um zu prüfen, ob die TPU-Slices miteinander kommunizieren können.

JAX auf jedem der Slices installieren – JAX-Code auf TPU-Pod-Slices ausführen

Behaupten Sie, dass len(jax.devices()) der Anzahl der Chips in Ihrer Multislice-Umgebung entspricht. Führen Sie dazu für jedes Segment den folgenden Befehl aus:

export JAX_USE_PJRT_C_API_ON_TPU=1 && python3 -c 'import jax; print(jax.devices())'

Wenn Sie diesen Code auf vier Slices von v4-16 ausführen, gibt es acht Chips pro Segment und vier Segmente. Insgesamt sollten also von jax.devices() insgesamt 32 Chips (Geräte) zurückgegeben werden.

Ressourcen in der Warteschlange auflisten

Sie können den Status der QR-Codes mit dem Befehl queued-resources list abrufen:

$ gcloud alpha compute tpus queued-resources list

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

Ressourcen in der Warteschlange beschreiben

Wenn Sie eine detaillierte Konfiguration und den Status eines QR-Codes sehen möchten, verwenden Sie die QR-API „describe“. Sie können diese API mit gcloud oder curl aufrufen.

mit gcloud:

$ gcloud alpha compute tpus queued-resources describe ${QUEUED_RESOURCE_ID}
...state:
 state: ACTIVE
...

mit curl:

$ curl -X GET -H "Authorization: Bearer $(gcloud auth print-access-token)" -H "Content-Type: application/json" https://tpu.googleapis.com/v2alpha1/projects/${PROJECT}/locations/${ZONE}/queuedResources/${QUEUED_RESOURCE_ID}
{
  "name": your-queued-res,
  "tpu": {
    "nodeSpec": [
      {
        ... // node 1
      },
      {
        ... // node 2
      },
      ...
    ]
  },
  ...
  "state": "ACTIVE"
  }
}

"state" stellt den Status eines QR-Codes dar. Die Zustände sind folgendermaßen definiert:

ACCEPTED (eingestellt, ersetzt durch WAITING_FOR_RESOURCES)	Der QR-create-Befehl wurde empfangen und beginnt mit der Bereitstellung, sobald Ressourcen verfügbar sind.
PROVISIONING	Die TPU-Slices werden bereitgestellt.
ACTIVE	Alle TPUs sind bereitgestellt und einsatzbereit. Wenn ein Startskript angegeben wird, wird es auf allen TPUs ausgeführt, wenn der QR-Status zu ACTIVE wechselt. Bei einem kleinen Programmfehler wird derzeit der QR-Status als ACTIVE angezeigt, während einige Segmente gelöscht werden. In diesem Fall funktioniert der QR-Code nicht wie erwartet und alle Ressourcen sollten bereinigt werden.
FAILED	Die Segmente konnten nicht bereitgestellt werden.
SUSPENDING	Ein oder mehrere Segmente werden gelöscht.
SUSPENDED	Alle zugrunde liegenden Segmente werden gelöscht, aber der QR-Code bleibt erhalten, bis er explizit gelöscht wird. Derzeit kann ein gesperrter QR-Code nicht fortgesetzt werden und sollte gelöscht werden.
DELETING	Der QR-Code wird gelöscht.

Job in einer bereitgestellten Umgebung starten

Sie können Arbeitslasten manuell ausführen, indem Sie über SSH eine Verbindung zu allen Hosts in jedem Segment herstellen und auf allen Hosts den folgenden Befehl ausführen. In der Regel müssen Sie für jedes Segment einmal einen Befehl wie den folgenden ausführen.

$ gcloud compute tpus tpu-vm ssh ${QUEUED_RESOURCE_ID}-0 \
  --zone ${ZONE} \
  --worker=all \
  --command="command-to-run"

Ressourcen in der Warteschlange werden zurückgesetzt

Mit der ResetQueuedResource API können alle VMs in einem ACTIVE-QR zurückgesetzt werden. Durch das Zurücksetzen der VMs wird der Arbeitsspeicher der Maschine zwangsweise gelöscht und die VM in ihren ursprünglichen Zustand zurückgesetzt. Alle lokal gespeicherten Daten bleiben erhalten und das Startskript wird nach dem Zurücksetzen aufgerufen. Die ResetQueuedResource API kann hilfreich sein, wenn Sie alle TPUs neu starten möchten. Wenn beispielsweise das Training hängen geblieben ist und das Zurücksetzen aller VMs einfacher ist als die Fehlerbehebung.

Das Zurücksetzen aller VMs wird parallel durchgeführt. Ein ResetQueuedResource-Vorgang dauert 1 bis 2 Minuten. Verwenden Sie den folgenden Befehl, um die API aufzurufen:

$ gcloud alpha compute tpus queued-resources reset $QUEUED_RESOURCE_ID

Ressourcen in der Warteschlange löschen

Wenn Sie Ressourcen am Ende Ihrer Trainingseinheit freigeben möchten, löschen Sie die Warteschlange in der Warteschlange mit dem Flag --force. Das Löschen dauert zwei bis fünf Minuten und kann im Hintergrund mit dem optionalen Flag --async ausgeführt werden.

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

Automatische Wiederherstellung nach Fehlern

Bei einer Störung bietet Multislice eine eingriffsfreie Reparatur des betroffenen Segments und anschließend das Zurücksetzen aller Segmente. Das betroffene Segment wird durch ein neues Segment ersetzt und die übrigen ansonsten fehlerfreien Segmente werden zurückgesetzt. Wenn keine Kapazität zum Zuweisen eines Ersatzsegments verfügbar ist, wird das Training beendet.

Damit das Training nach einer Unterbrechung automatisch fortgesetzt wird, müssen Sie ein Startskript angeben, das die zuletzt gespeicherten Prüfpunkte prüft und lädt. Das Startskript wird automatisch jedes Mal ausgeführt, wenn ein Segment neu zugewiesen oder eine VM zurückgesetzt wird. Sie geben in der JSON-Nutzlast, die Sie an die API zum Erstellen von QR-Anfragen senden, ein Startskript an.

Mit dem folgenden Startskript, das in Ressourcen in der Warteschlange erstellt verwendet wird, können Sie nach Fehlern automatisch wiederherstellen und das Training über Prüfpunkte fortsetzen, die während des MaxText-Trainings in einem Cloud Storage-Bucket gespeichert sind:

{
 "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 es ausprobieren.

Profilerstellung und Fehlerbehebung

Die Profilerstellung ist für einzelne Slice- und Multislice-Umgebungen identisch. Weitere Informationen finden Sie unter Profilerstellung für JAX-Programme.

Training optimieren

Fragmentierung mit Multislice für maximale Leistung

Um in Multislice-Umgebungen für eine maximale Leistung zu sorgen, müssen Sie überlegen, wie Fragmente in mehrere Segmente aufgeteilt werden sollen. In der Regel gibt es drei Optionen: Datenparallelität, vollständig fragmentierte Datenparallelität und Pipeline-Parallelität. Wir raten davon ab, Fragmentierungen für die Modelldimensionen (manchmal als Tensor-Parallelität) zu aktivieren, da sie zu viel Interslic-Band erfordern. Bei all diesen Strategien können Sie dieselbe Fragmentierungsstrategie in einem Segment beibehalten, das in der Vergangenheit für Sie funktioniert hat.

Wir empfehlen, mit einer reinen Datenparallelität zu beginnen. Die Verwendung einer vollständig fragmentierten Datenparallelität ist hilfreich, um die Arbeitsspeichernutzung freizugeben. Der Nachteil ist, dass die Kommunikation zwischen Segmenten das DCN-Netzwerk nutzt und Ihre Arbeitslast verlangsamt. Verwenden Sie die Pipeline-Parallelität nur bei Bedarf anhand der Batchgröße (wie unten analysiert).

Wann ist Datenparallelität zu verwenden?

Reine Datenparallelität funktioniert gut, wenn Sie eine Arbeitslast haben, die gut läuft, Sie ihre Leistung aber verbessern möchten, indem Sie mehrere Segmente skalieren.

Für eine starke Skalierung über mehrere Segmente hinweg muss die Zeit, die für die vollständige Reduzierung über DCN benötigt wird, kürzer als die für die Ausführung eines Rückwärtsdurchlaufs benötigte Zeit sein. DCN wird für die Kommunikation zwischen Segmenten verwendet und ist ein begrenzender Faktor für den Arbeitslastdurchsatz.

Jeder v4-TPU-Chip erreicht einen Spitzenwert von 275 × 10¹² FLOPS/Sekunde. Es gibt vier Chips pro TPU-Host und jeder Host hat eine maximale Netzwerkbandbreite von 50 Gbit/s. Das bedeutet, dass die arithmetische Intensität 4 × 275 × 10¹² FLOPS / 50 Gbit/s = 22.000 FLOPS/Bit ist.

Ihr Modell verwendet für jeden Parameter pro Schritt 32 bis 64 Bits der DCN-Bandbreite. Wenn Sie zwei Segmente verwenden, nutzt Ihr Modell 32 Bit DCN-Bandbreite. Wenn Sie mehr als zwei Segmente verwenden, führt der Compiler einen Voll-Shuffle-Vorgang mit Reduzierung aus. Für jeden Parameter verwenden Sie bis zu 64 Bits für die DCN-Bandbreite.

Die für jeden Parameter erforderliche FLOPS-Menge variiert je nach Modell.

Insbesondere bei Transformatoren basierend auf Sprachmodellen beträgt die Anzahl der erforderlichen FLOPS etwa 6 * B * P.

• 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 umgekehrten Vorgangs beträgt 4 * B.

Damit eine starke Skalierung über mehrere Segmente hinweg erzielt werden kann, muss die operative Intensität die arithmetische Intensität der TPU-Hardware überschreiten. Zur Berechnung der operativen Intensität teilen Sie die Anzahl der FLOPS pro Parameter während des Rückwärtsdurchlaufs durch die Netzwerkbandbreite (in Bit) pro Parameter pro Schritt:

Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

Wenn Sie ein Skelettmodell verwenden und zwei Segmente verwenden, gilt Folgendes:

Operational intensity = 4 * B  / 32

Wenn Sie mehr als zwei Segmente verwenden, ist Folgendes zu beachten:

Operational intensity = 4 * B  / 64

Dies deutet auf eine Mindest-Batchgröße zwischen 176.000 und 352.000 für Transformer-basierte Sprachmodelle hin. Da das DCN-Netzwerk Pakete kurzzeitig löschen kann, ist es am besten, eine erhebliche Fehlertoleranz beizubehalten und die Datenparallelität nur dann bereitzustellen, wenn die Batchgröße pro Pod mindestens 350.000 (2 Pods) bis 700.000 (zahlreiche Pods) beträgt.

Bei anderen Modellarchitekturen müssen Sie die Laufzeit Ihrer Rücklaufkarte pro Segment schätzen (entweder durch Timing mit einem Profiler oder durch Zählen von FLOPS). Anschließend können Sie diese mit der erwarteten Laufzeit vergleichen, um die DCN zu reduzieren und eine gute Schätzung zu treffen, ob eine Datenparallelität für Sie sinnvoll ist.

Wann wird vollständig fragmentierte Datenparallelität (FSDP) verwendet?

Vollständig fragmentierte Datenparallelität (FSDP) kombiniert Datenparallelität (Fragmentierung der Daten über Knoten hinweg) mit Fragmentierung der Gewichtungen über alle Knoten. Für jeden Vorgang in den Vorwärts- und Rückwärtsläufen werden die Gewichtungen vollständig erfasst, sodass jedes Segment die benötigte Gewichtung hat. Statt die Farbverläufe mit der Reduzierung des Farbverlaufs zu synchronisieren, werden die Gradienten bei der Produktion reduziert. Auf diese Weise erhält jedes Segment nur die Farbverläufe für die Gewichtungen, für die es verantwortlich ist.

Ähnlich wie bei der Datenparallelität muss FSDP die globale Batchgröße linear mit der Anzahl der Segmente skalieren. FSDP senkt den Speicherdruck, wenn Sie die Anzahl der Segmente erhöhen. Dies liegt daran, dass die Anzahl der Gewichtungen und der Status des Optimierungstools pro Segment reduziert wird, dafür aber der Netzwerkdatenverkehr und die Blockierungswahrscheinlichkeit aufgrund eines verzögerten Kollektivs höher sind.

In der Praxis ist es am besten, wenn Sie den FSDP für die Segmente nutzen, indem Sie den Batch pro Segment erhöhen. So können Sie mehr Aktivierungen speichern, um die Rematerialisierung während der Zurückleitung zu minimieren oder die Anzahl der Parameter in Ihrem neuronalen Netzwerk zu erhöhen.

Die Vorgänge mit vollständiger Erfassung und der reduzierten Reduzierung funktionieren in FSDP ähnlich wie in DP. Sie können also ermitteln, ob Ihre FSDP-Arbeitslast auf die gleiche Weise durch die DCN-Leistung eingeschränkt ist, wie im vorherigen Abschnitt beschrieben.

Wann die Pipeline-Parallelität verwendet werden sollte

Die Pipeline-Parallelität wird relevant, wenn Sie eine hohe Leistung mit anderen Parallelisierungsstrategien erzielen, die eine globale Batchgröße erfordern, die größer als Ihre bevorzugte maximale Batchgröße ist. Mit der Pipeline-Parallelität können die Segmente aus einer Pipeline einen Batch „teilen“. Die Pipeline-Parallelität hat jedoch zwei wichtige Nachteile:

1. Es tritt die „Pipeline-Bubble“ auf, in der Chips inaktiv sind, weil sie auf Daten warten.
2. Er erfordert Mikro-Batching, wodurch die effektive Batchgröße und arithmetische Intensität reduziert und schließlich die FLOP-Auslastung modelliert wird.

Die Pipeline-Parallelität sollte nur verwendet werden, wenn die anderen Parallelisierungsstrategien eine zu große globale Batchgröße erfordern. Bevor Sie die Pipeline-Parallelität ausprobieren, lohnt es sich, zu testen, ob sich die Konvergenz pro Stichprobe bei der Batchgröße verlangsamt, die erforderlich ist, um einen leistungsfähigen FSDP zu erreichen. FSDP neigt dazu, eine höhere FLOP-Modellmodellierung zu erzielen, aber wenn die Konvergenz pro Stichprobe langsam wird, wenn die Batchgröße zunimmt, ist die Pipeline-Parallelität möglicherweise noch die bessere Wahl. Die meisten Arbeitslasten können ausreichend große Batchgrößen tolerieren, um nicht von der Pipeline-Parallelität zu profitieren. Ihre Arbeitslast kann sich jedoch unterscheiden.

Wenn die Pipeline-Parallelität erforderlich ist, empfehlen wir, sie mit Datenparallelität oder FSDP zu kombinieren. Auf diese Weise können Sie die Pipelinetiefe minimieren und gleichzeitig die Batchgröße pro Pipeline erhöhen, bis die DCN-Latenz den Durchsatz eines Teils verringert. Wenn Sie also N-Segmente haben, sollten Sie Pipelines der Tiefe 2 und N/2 Replikate der Datenparallelität und dann Pipelines der Tiefe 4 und N/4-Replikate der Datenparallelität so lange betrachten, bis der Batch pro Pipeline so groß wird, dass die DCN-Sammlungen hinter der Arithmetik im Rückwärtspass verborgen sein können. Dadurch wird die durch die Pipeline-Parallelität verursachte Verlangsamung minimiert, während Sie das globale Batchgrößenlimit weiter skalieren können.

Best Practices für die Multiplikation

Laden der Daten

Während des Trainings laden wir wiederholt Batches aus einem Dataset, um sie in das Modell zu importieren. Ein effizienter, asynchroner Datenlader, der den Batch über mehrere Hosts fragmentiert, ist wichtig, damit die TPUs der Arbeit nicht ausgelastet werden. Beim aktuellen Data Loader in MaxText hat jeder Host die gleiche Teilmenge der Beispiele. Diese Lösung ist für Text geeignet, erfordert jedoch eine Reshard im Modell. Außerdem bietet MaxText noch keinen deterministischen Snapshot, mit dem der Daten-Iterator dieselben Daten vor und nach dem vorzeitigen Beenden laden könnte.

Prüfpunktausführung

Die Orbax-Checkpoint-Bibliothek bietet Primitive zum Prüfen von JAX PyTrees im lokalen Speicher oder in Google Cloud Storage. Eine Referenzintegration mit synchroner Prüfpunktprüfung in MaxText wird in checkpointing.py bereitgestellt.

Training optimieren

Fragmentierung mit Multislice für maximale Leistung

Um in Multislice-Umgebungen für eine maximale Leistung zu sorgen, müssen Sie überlegen, wie Fragmente in mehrere Segmente aufgeteilt werden sollen. In der Regel gibt es drei Optionen: Datenparallelität, vollständig fragmentierte Datenparallelität und Pipeline-Parallelität. Wir raten davon ab, Fragmentierungen für die Modelldimensionen (manchmal als Tensor-Parallelität) zu aktivieren, da sie zu viel Interslic-Band erfordern. Bei all diesen Strategien können Sie dieselbe Fragmentierungsstrategie in einem Segment beibehalten, das in der Vergangenheit für Sie funktioniert hat.

Wir empfehlen, mit einer reinen Datenparallelität zu beginnen. Die Verwendung einer vollständig fragmentierten Datenparallelität ist hilfreich, um die Arbeitsspeichernutzung freizugeben. Der Nachteil ist, dass die Kommunikation zwischen Segmenten das DCN-Netzwerk nutzt und Ihre Arbeitslast verlangsamt. Verwenden Sie die Pipeline-Parallelität nur bei Bedarf basierend auf der Batchgröße (wie unten analysiert).

Wann sollte ich Datenparallelität verwenden?

Reine Datenparallelität funktioniert gut, wenn Sie eine Arbeitslast haben, die gut läuft, Sie ihre Leistung aber verbessern möchten, indem Sie mehrere Segmente skalieren.

Für eine starke Skalierung über mehrere Segmente hinweg muss die Zeit, die für die vollständige Reduzierung über DCN benötigt wird, kürzer als die für die Ausführung eines Rückwärtsdurchlaufs benötigte Zeit sein. DCN wird für die Kommunikation zwischen Segmenten verwendet und ist ein begrenzender Faktor für den Arbeitslastdurchsatz.

Jeder v4-TPU-Chip erreicht einen Spitzenwert von 275 × 10¹² FLOPS/Sekunde. Es gibt vier Chips pro TPU-Host und jeder Host hat eine maximale Netzwerkbandbreite von 50 Gbit/s. Das bedeutet, dass die arithmetische Intensität 4 × 275 × 10¹² FLOPS / 50 Gbit/s = 22.000 FLOPS/Bit ist.

Ihr Modell verwendet für jeden Parameter pro Schritt 32 bis 64 Bits der DCN-Bandbreite. Wenn Sie zwei Segmente verwenden, nutzt Ihr Modell 32 Bit DCN-Bandbreite. Wenn Sie mehr als zwei Segmente verwenden, führt der Compiler einen Voll-Shuffle-Vorgang mit Reduzierung aus. Für jeden Parameter verwenden Sie bis zu 64 Bits für die DCN-Bandbreite.

Die für jeden Parameter erforderliche FLOPS-Menge variiert je nach Modell.

Insbesondere bei Transformatoren basierend auf Sprachmodellen beträgt die Anzahl der erforderlichen FLOPS etwa 6 * B * P.

• 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 umgekehrten Vorgangs beträgt 4 * B.

Damit eine starke Skalierung über mehrere Segmente hinweg erzielt werden kann, muss die operative Intensität die arithmetische Intensität der TPU-Hardware überschreiten. Wir haben die arithmetische Intensität zu Beginn dieses Abschnitts berechnet. Zur Berechnung der operativen Intensität teilen wir die Anzahl der FLOPS pro Parameter während des Rückwärtslaufs durch die Netzwerkbandbreite (in Bit) pro Parameter pro Schritt:

Operational Intensity = FLOPSbackwards_pass / DCN bandwidth

Wenn wir für ein Transformer-basiertes Sprachmodell zwei Segmente verwenden, gilt Folgendes:

Operational intensity = 4 * B  / 32

Wenn wir mehr als zwei Segmente verwenden, ist Folgendes zu beachten:

Operational intensity = 4 * B  / 64

Dies deutet auf eine Mindest-Batchgröße zwischen 176.000 und 352.000 für Transformer-basierte Sprachmodelle hin. Da das DCN-Netzwerk Pakete kurzzeitig löschen kann, ist es am besten, eine erhebliche Fehlertoleranz beizubehalten und die Datenparallelität nur dann bereitzustellen, wenn die Batchgröße pro Pod mindestens 350.000 (zwei Pods) bis 700.000 (zahlreiche Pods) beträgt.

Bei anderen Modellarchitekturen müssen Sie die Laufzeit Ihrer Rücklaufkarte pro Segment schätzen (entweder durch Timing mit einem Profiler oder durch Zählen von FLOPS). Anschließend können Sie diese mit der erwarteten Laufzeit vergleichen, um die DCN zu reduzieren und eine gute Schätzung zu treffen, ob eine Datenparallelität für Sie sinnvoll ist.

Unterstützte Konfigurationen

Formen

Alle Segmente müssen die gleiche Form haben (z. B. dasselbe AcceleratorType). Heterogene Segmentformen werden derzeit nicht unterstützt.

TPU-Generationen

Die Multislice-Funktion wird derzeit nur für v4-Pods unterstützt.

Orchestrierung

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

Frameworks

Die Multislice-Vorschau in v4 unterstützt nur JAX-Arbeitslasten.

Parallelität

Wir empfehlen Nutzern, Multislice mit Datenparallelität in dieser Vorschau zu testen. Kontaktieren Sie uns, wenn Sie mehr über die Implementierung der Pipeline-Parallelität mit Multislice erfahren möchten.

Support und Feedback

Wir freuen uns über jedes Feedback! Wenn Sie Feedback geben oder Support anfordern möchten, kontaktieren Sie uns oder senden Sie eine E-Mail an [e-mail].

Nutzungsbedingungen

Alle Informationen, die Google Ihnen zu dieser privaten Vorschau bereitgestellt hat, sind vertrauliche Informationen von Google und unterliegen den Vertraulichkeitsbestimmungen der Nutzungsbedingungen der Google Cloud Platform oder einer anderen Vereinbarung, die Ihre Nutzung von Google Cloud regelt.