Cloud TPU Multislice

Cloud TPU Multislice ist eine Full-Stack-Technologie zur Leistungsskalierung, mit der ein Trainingsjob mehrere TPU-Slices innerhalb eines einzelnen Slice oder auf Slices in mehreren Pods mit standardmäßiger Datenparallelität verwenden kann. Mit TPU v4-Chips können für Trainingsjobs in einer einzelnen Ausführung mehr als 4.096 Chips verwendet werden. Für Trainingsjobs, für die weniger als 4.096 Chips erforderlich sind, kann ein einzelner Slice die beste Leistung bieten. Mehrere kleinere Slices sind jedoch leichter verfügbar, was einen schnelleren Start ermöglicht, wenn Multislice mit kleineren Slices verwendet wird.

Bei der Bereitstellung in Multislice-Konfigurationen kommunizieren TPU-Chips in jedem Slice über Inter-Chip-Interconnect (ICI). TPU-Chips in verschiedenen Slices kommunizieren, indem sie Daten an CPUs (Hosts) übertragen, die die Daten wiederum über das Rechenzentrumsnetzwerk (Data Center Network, DCN) übertragen. Weitere Informationen zur Skalierung mit Multislice finden Sie im Blogpost How to scale AI training to up to tens of thousands of Cloud TPU chips with Multislice.

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

Konzepte

Beschleunigertyp

Die Form der TPU-Slices, aus denen ein Multislice besteht. Jeder Slice in einer Multislice-Anfrage hat denselben Beschleunigertyp. Ein Beschleunigertyp besteht aus einem TPU-Typ (v4 oder höher) gefolgt von der Anzahl der TensorCores. Beispiel: v5litepod-128 gibt eine v5e-TPU mit 128 TensorCores an.

Automatische Reparatur

Wenn bei einem Slice ein Wartungsereignis, ein vorzeitiges Beenden oder ein Hardwarefehler auftritt, wird von Cloud TPU ein neuer Slice erstellt. Wenn nicht genügend Ressourcen zum Erstellen eines neuen Slice vorhanden sind, wird die Erstellung erst abgeschlossen, wenn Hardware verfügbar ist. Nachdem der neue Slice erstellt wurde, werden alle anderen Slices in der Multislice-Umgebung neu gestartet, damit das Training fortgesetzt werden kann. Mit einem richtig konfigurierten Startscript kann das Trainingsscript automatisch ohne Eingreifen des Nutzers neu gestartet werden. Dabei wird der letzte Prüfpunkt geladen und das Training wird fortgesetzt.

Rechenzentrumsnetzwerk (Data Center Network, DCN)

Das DCN ist ein Netzwerk mit höherer Latenz und geringerem Durchsatz (im Vergleich zu ICI) und verbindet TPU-Slices in einer Multislice-Konfiguration.

Planung in Gruppen

Wenn alle TPU-Slices gemeinsam bereitgestellt werden, werden definitiv entweder alle oder keine der Slices erfolgreich bereitgestellt.

Inter-Chip-Interconnect (ICI)

Interne Links mit hoher Geschwindigkeit und geringer Latenz verbinden TPUs in einem TPU-Pod.

Multislice

Der Begriff wird für zwei oder mehr TPU-Chip-Slices verwendet, die über das DCN kommunizieren können.

Knoten

Im Kontext von Multislice bezieht sich der Begriff „Knoten“ auf einen einzelnen TPU-Slice. Jedem TPU-Slice in einem Multislice wird eine Knoten-ID zugewiesen.

Startscript

Ein standardmäßiges Compute Engine-Startscript, das jedes Mal ausgeführt wird, wenn eine VM gestartet oder neu gestartet wird. Bei Multislice wird sie in der Anfrage zur Erstellung von Ressourcen für die Warteschlange angegeben. Weitere Informationen zu Cloud TPU-Startscripts finden Sie unter TPU-Ressourcen verwalten.

Tensor

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

Cloud TPU-Kapazitätstypen

TPUs können mit verschiedenen Kapazitätstypen erstellt werden (siehe „Nutzungsoptionen“ auf der Seite für Cloud TPU-Preise):

• Reservierung: Um eine Reservierung nutzen zu können, benötigen Sie eine Reservierungsvereinbarung mit Google. Verwenden Sie beim Erstellen der Ressourcen das Flag --reserved.

• Spot: Es werden Spot-VMs verwendet, die auf ein Kontingent auf Abruf angerechnet werden. Ihre Ressourcen können vorzeitig beendet werden, um Platz für Anfragen für einen Job mit höherer Priorität zu schaffen. Verwenden Sie beim Erstellen von Ressourcen das Flag --spot.

• On-Demand: Ein On-Demand-Kontingent, für das keine Reservierung erforderlich ist und bei dem Ressourcen nicht vorzeitig beendet werden. Die TPU-Anfrage wird in eine Warteschlange für ein On-Demand-Kontingent von Cloud TPU gestellt. Die Verfügbarkeit von Ressourcen kann nicht gewährleistet werden. Standardmäßig ausgewählt, keine Flags erforderlich.

Jetzt starten

1. Richten Sie die Cloud TPU-Umgebung ein.

2. In the Google Cloud console, activate Cloud Shell.

   Activate Cloud Shell

   At the bottom of the Google Cloud console, a Cloud Shell session starts and displays a command-line prompt. Cloud Shell is a shell environment with the Google Cloud CLI already installed and with values already set for your current project. It can take a few seconds for the session to initialize.

Wenn Sie Multislice verwenden möchten, müssen Ihre TPU-Ressourcen als in die Warteschlange gestellte Ressourcen verwaltet werden.

Einführendes Beispiel

In dieser Anleitung wird Code aus dem MaxText-GitHub-Repository verwendet. MaxText ist ein leistungsstarkes, beliebig skalierbares, einfaches Open-Source-LLM, das in Python und JAX geschrieben wurde und gut getestet ist. MaxText wurde für effizientes Training auf Cloud TPU entwickelt.

Der Code in shardings.py soll Ihnen den Einstieg in die verschiedenen Parallelitätsoptionen erleichtern. Dazu gehören beispielsweise Datenparallelität, vollständig fragmentierte Datenparallelität (Fully Sharded Data Parallelism, FSDP) und Tensorparallelität. Der Code lässt sich von einzelnen Slices auf Multislice-Umgebungen skalieren.

ICI-Parallelität

ICI bezieht sich auf die Hochgeschwindigkeitsverbindung, die die TPUs in einem Slice verbindet. ICI-Fragmentierung entspricht der Fragmentierung innerhalb eines Slice. shardings.py bietet drei Parameter für ICI-Parallelität:

• 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 werden, dass ici_data_parallelism * ici_fsdp_parallelism * ici_tensor_parallelism der Anzahl der Chips im Slice entspricht.

In der folgenden Tabelle finden Sie Beispiele für Nutzereingaben für die ICI-Parallelität für die vier in v4-8 verfügbaren Chips:

	ici_data_parallelism	ici_fsdp_parallelism	ici_tensor_parallelism
4-Wege-FSDP	1	4	1
4-Wege-Tensorparallelität	1	1	4
2-Wege-FSDP + 2-Wege-Tensorparallelität	1	2	2

ici_data_parallelism sollte in den meisten Fällen auf 1 belassen werden, da das ICI-Netzwerk schnell genug ist, um FSDP fast immer der Datenparallelität vorzuziehen.

In diesem Beispiel wird davon ausgegangen, dass Sie mit dem Ausführen von Code auf einem einzelnen TPU-Slice vertraut sind, zum Beispiel wie unter Berechnung mit JAX auf einer Cloud TPU-VM ausführen beschrieben. In diesem Beispiel wird gezeigt, wie Sie shardings.py für einen Slice ausführen.

1. Richten Sie die Umgebung ein:

   $ gcloud auth login
   $ export QR_ID=your-queued-resource-id
   $ export TPU_NAME=your-tpu-name
   $ export PROJECT=your-project-name
   $ export ZONE=us-central1-a
   $ export NETWORK_NAME=your-network-name
   $ export SUBNETWORK_NAME=your-subnetwork-name
   $ export RUNTIME_VERSION=v2-alpha-tpuv5-lite
   $ export ACCELERATOR_TYPE=v5litepod-16
   $ export EXAMPLE_TAG_1=your-tag-1
   $ export EXAMPLE_TAG_2=your-tag-2
   $ export SLICE_COUNT=4
   $ export STARTUP_SCRIPT='#!/bin/bash\n'

   Variablenbeschreibungen

   Eingabe	Beschreibung
   QR_ID	Vom Nutzer zugewiesene ID der Anfrage für in die Warteschlange gestellte Ressourcen
   TPU_NAME	Vom Nutzer zugewiesener TPU-Name
   PROJEKT	Projektname inGoogle Cloud
   ZONE	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 Version derCloud TPU-Software
   ACCELERATOR_TYPE	v4-16
   EXAMPLE_TAG_1, EXAMPLE_TAG_2 …	Tags, mit denen gültige Quellen oder Ziele für Netzwerkfirewalls identifiziert werden
   SLICE_COUNT	Anzahl der Slices. Auf maximal 256 Slices beschränkt.
   STARTUP_SCRIPT	Wenn Sie ein Startscript angeben, wird es ausgeführt, wenn der TPU-Slice bereitgestellt oder neu gestartet wird.

2. Erstellen Sie SSH-Schlüssel für gcloud. Wir empfehlen, das Passwort leer zu lassen (drücken Sie nach dem Ausführen des folgenden Befehls zweimal die Eingabetaste). Wenn angezeigt wird, dass 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 bereit:

   gcloud

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --node-id=${TPU_NAME} \
       --zone=${ZONE} \
       [--reserved |--spot]

   Die Google Cloud CLI unterstützt nicht alle Optionen zum Erstellen von Ressourcen für die Warteschlange, z. B. Tags. Weitere Informationen finden Sie unter Ressourcen für die Warteschlange erstellen.

   Console

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

      Zu TPUs

   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 Sie die TPU erstellen möchten.

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

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

   7. Klicken Sie auf die Umschaltfläche Warteschlange aktivieren.

   8. Geben Sie im Feld Name der Ressource in der Warteschlange einen Namen für die Anfrage für in die Warteschlange gestellte Ressourcen ein.

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

4. Warten Sie, bis sich die Ressource in der Warteschlange im Status ACTIVE befindet. Das bedeutet, dass sich die Worker-Knoten im Status READY befinden. Wenn die Bereitstellung der Ressource aus der Warteschlange beginnt, kann es je nach Größe der Ressource ein bis fünf Minuten dauern, bis sie abgeschlossen ist. Sie können den Status einer Anfrage für in die Warteschlange gestellte Ressourcen mit der gcloud CLI oder der Console von Google Cloud prüfen:

   gcloud

   $ gcloud compute tpus queued-resources \
       list --filter=${QR_ID} --zone=${ZONE}

   Console

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

      Zu TPUs

   2. Klicken Sie auf den Tab In die Warteschlange gestellte Ressourcen.

   3. Klicken Sie auf den Namen Ihrer Anfrage für in die Warteschlange gestellte Ressourcen.

5. Stellen Sie eine SSH-Verbindung zur TPU-VM her:

   $ gcloud compute tpus tpu-vm ssh ${TPU_NAME} --zone=${ZONE}

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

   $ git clone https://github.com/AI-Hypercomputer/maxtext && cd maxtext

7. Installieren Sie Python 3.10:

   $ sudo apt-get update
   $ sudo apt install python3.10
   $ sudo apt install python3.10-venv

8. Erstellen und aktivieren Sie eine virtuelle Umgebung:

   $ python3 -m venv your-venv-name
   $ source your-venv-name/bin/activate

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

   $ bash setup.sh

10. Führen Sie den folgenden Befehl aus, um shardings.py für Ihren TPU-Slice auszuführen.

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

    Die Ergebnisse werden in den Logs angezeigt. Ihre TPUs sollten etwa 260 TFLOP pro Sekunde oder eine beeindruckende FLOP-Auslastung von über 90 % erreichen. In diesem Fall haben wir ungefähr den maximalen Batch ausgewählt, der in das HBM (High Bandwidth Memory) der TPU passt.

11. Sie können auch andere Fragmentierungsstrategien über ICI ausprobieren, z. B. die folgende Kombination:

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

12. Löschen Sie die in die Warteschlange gestellte Ressource und den TPU-Slice, wenn Sie fertig sind. Sie sollten diese Bereinigungsschritte in der Umgebung ausführen, in der Sie den Slice 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.

    gcloud

    $ gcloud compute tpus queued-resources \
        delete ${QR_ID} --force (--async)

    Console

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

       Zu TPUs

    2. Klicken Sie auf den Tab In die Warteschlange gestellte Ressourcen.

    3. Aktivieren Sie das Kästchen neben Ihrer Anfrage für in die Warteschlange gestellte Ressourcen.

    4. Klicken Sie auf delete Löschen.

Multislice-Fragmentierung mit DCN-Parallelität

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

• dcn_data_parallelism
• dcn_fsdp_parallelism
• dcn_tensor_parallelism

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

Beispiel für zwei Slices: --dcn_data_parallelism = 2

	dcn_data_parallelism	dcn_fsdp_parallelism	dcn_tensor_parallelism	Anzahl der Slices
2-Wege-Datenparallelität	2	1	1	2

dcn_tensor_parallelism sollte immer auf 1 festgelegt werden, da das DCN für eine solche Fragmentierung nicht gut geeignet ist. Bei typischen LLM-Arbeitslasten auf v4-Chips sollte dcn_fsdp_parallelism auch auf 1 gesetzt werden. Daher sollte dcn_data_parallelism auf die Anzahl der Slices gesetzt werden. Dies ist jedoch anwendungsabhängig.

Mit zunehmender Anzahl der Slices (vorausgesetzt, Slice-Größe und Batch pro Slice bleiben konstant) erhöht sich der Datenparallelitätsgrad.

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 sehr denen in Getting Started: Quick Experiments on Multiple Slices aus dem MaxText-Repository. Hier führen wir jedoch shardings.py anstelle des komplexeren LLM in train.py aus.

Das Tool multihost_runner.py ist für schnelle Experimente optimiert, bei denen dieselben TPUs wiederholt verwendet werden. Da das Script multihost_runner.py von langlebigen SSH-Verbindungen abhängt, empfehlen wir es nicht für Jobs mit langer Ausführungszeit. Wenn Sie einen längeren Job ausführen möchten (z. B. über Stunden oder Tage), empfehlen wir, multihost_job.py zu verwenden.

In dieser Anleitung wird der Begriff Runner für den Computer verwendet, auf dem Sie das Script multihost_runner.py ausführen. Der Begriff Worker wird für die TPU-VMs verwendet, aus denen Ihre Slices bestehen. Sie können multihost_runner.py auf einem lokalen Computer oder einer beliebigen Compute Engine-VM im selben Projekt wie Ihre Slices ausführen. Die Ausführung 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 führen Sie shardings.py auf zwei v5e-16-Slices aus, also auf insgesamt vier VMs und 16 TPU-Chips. Sie können das Beispiel für die Ausführung auf mehr TPUs anpassen.

Umgebung einrichten

1. Klonen Sie MaxText auf Ihrem Runner-Computer:

   $ git clone https://github.com/AI-Hypercomputer/maxtext

2. Wechseln Sie in das Repository-Verzeichnis.

   $ cd maxtext

3. Erstellen Sie SSH-Schlüssel für gcloud. Wir empfehlen, das Passwort leer zu lassen (drücken Sie nach dem Ausführen des folgenden Befehls zweimal die Eingabetaste). Wenn angezeigt wird, dass die Datei google_compute_engine bereits vorhanden ist, wählen Sie, die vorhandene Version nicht zu behalten.

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

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

     $ export SLICE_COUNT=2
     

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

   gcloud

   Der folgende Befehl zeigt, wie Sie eine v5e-Multislice-TPU erstellen. Wenn Sie eine andere TPU-Version verwenden möchten, geben Sie einen anderen accelerator-type und eine andere runtime-version an.

   $ gcloud compute tpus queued-resources \
       create ${QR_ID} \
       --accelerator-type=${ACCELERATOR_TYPE} \
       --runtime-version=${RUNTIME_VERSION} \
       --node-count=${SLICE_COUNT} \
       --node-prefix=${TPU_NAME} \
       --zone=${ZONE} \
       [--reserved|--spot]

   Console

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

      Zu TPUs

   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 Sie die TPU erstellen möchten.

   5. Wählen Sie im Feld TPU-Typ einen Beschleunigertyp aus. Der Beschleunigertyp gibt die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Multislice wird nur auf Cloud TPU v4 und späteren TPU-Versionen unterstützt. Weitere Informationen zu TPU-Versionen finden Sie unter TPU-Versionen.

   6. Wählen Sie im Feld Version der TPU-Software eine Softwareversion aus. Beim Erstellen einer Cloud TPU-VM gibt das Feld „Version der TPU-Software“ die Version der auf den TPU-VMs zu installierenden TPU-Laufzeit an. Weitere Informationen finden Sie unter TPU-Softwareversionen.

   7. Klicken Sie auf die Umschaltfläche Warteschlange aktivieren.

   8. Geben Sie im Feld Name der Ressource in der Warteschlange einen Namen für die Anfrage für in die Warteschlange gestellte Ressourcen ein.

   9. Aktivieren Sie das Kästchen Multislice-TPU erstellen.

   10. Geben Sie im Feld Slice-Anzahl die Anzahl der Slices ein, die Sie erstellen möchten.

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

6. Wenn die Bereitstellung der Ressource aus der Warteschlange beginnt, kann es je nach Größe der Ressource ein bis fünf Minuten dauern, bis sie abgeschlossen ist. Warten Sie, bis sich die Ressource aus der Warteschlange im Status ACTIVE befindet. Sie können den Status einer Anfrage für in die Warteschlange gestellte Ressourcen mit der gcloud CLI oder der Console von Google Cloud prüfen:

   gcloud

   $ gcloud compute tpus queued-resources list \
       --filter=${QR_ID} --zone=${ZONE} --project=${PROJECT}

   Die generierte Ausgabe sollte in etwa so aussehen:

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

   Console

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

      Zu TPUs

   2. Klicken Sie auf den Tab In die Warteschlange gestellte Ressourcen.

   3. Klicken Sie auf den Namen Ihrer Anfrage für in die Warteschlange gestellte Ressourcen.

   Wenden Sie sich an Ihren Kundenbetreuer bei Google Cloud , wenn der Status der Ressource in der Warteschlange länger als 15 Minuten WAITING_FOR_RESOURCES oder PROVISIONING ist.

7. Installieren Sie die Abhängigkeiten:

   $ python3 multihost_runner.py \
       --TPU_PREFIX=${TPU_NAME} \
       --ZONE=${ZONE} \
       --COMMAND="bash setup.sh"

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

   $ python3 multihost_runner.py \
       --TPU_PREFIX=${TPU_NAME} \
       --ZONE=${ZONE} \
       --COMMAND="python3 -m pedagogical_examples.shardings \
       --dcn_data_parallelism ${SLICE_COUNT} \
       --ici_fsdp_parallelism 16 \
       --batch_size 131072 \
       --embedding_dimension 2048"

   In den Logdateien werden etwa 230 TFLOPS Leistung angezeigt.

   Weitere Informationen zum Konfigurieren der Parallelität finden Sie unter Multislice-Fragmentierung mit DCN-Parallelität und shardings.py.

9. Löschen 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 folgende Änderung am Code 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 sollte die einzige erforderliche Codeänderung bei der Umstellung auf Multislice sein. Um eine hohe Leistung zu erzielen, muss das DCN datenparallelen, vollständig fragmentierten datenparallelen oder pipelineparallelen Achsen zugeordnet werden. Aspekte der Leistung und Fragmentierungsstrategien werden ausführlicher unter Fragmentierung mit Multislice für maximale Leistung behandelt.

Um zu prüfen, ob Ihr Code auf alle Geräte zugreifen kann, können Sie nachsehen, ob len(jax.devices()) der Anzahl der Chips in Ihrer Multislice-Umgebung entspricht. Wenn Sie beispielsweise vier v4-16-Slices verwenden, haben Sie acht Chips pro Slice * vier Slices. len(jax.devices()) sollte also 32 zurückgeben.

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

Um eine lineare Beschleunigung zu erzielen, fügen Sie neue Slices mit derselben Größe wie Ihr vorhandener Slice hinzu. Wenn Sie beispielsweise einen v4-512-Slice verwenden, wird mit Multislice durch Hinzufügen eines zweiten v4-512-Slice und Verdoppeln der globalen Batchgröße etwa die doppelte Leistung erzielt. Weitere Informationen finden Sie unter Fragmentierung mit Multislice für maximale Leistung.

Job auf mehreren Slices ausführen

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

1. mit dem Script multihost_runner.py für experimentelle Runner
2. mit dem Script multihost_job.py für Produktions-Runner
3. manuell

Script für experimentelle Runner

Das Script multihost_runner.py verteilt Code auf eine vorhandene Multislice-Umgebung, führt Ihren Befehl auf jedem Host aus, kopiert Ihre Logs zurück und verfolgt den Fehlerstatus jedes Befehls. Das Script multihost_runner.py ist in der README-Datei für MaxText dokumentiert.

Da multihost_runner.py persistente SSH-Verbindungen aufrechterhält, eignet es sich nur für Experimente mit geringem Umfang und relativ kurzer Laufzeit. Sie können die Schritte in der Anleitung zu multihost_runner.py an Ihre Arbeitslast und Hardwarekonfiguration anpassen.

Script für Produktions-Runner

Bei Produktionsjobs, die eine hohe Ausfallsicherheit gegenüber Hardwarefehlern und anderen Unterbrechungen benötigen, empfiehlt es sich, die API zum Erstellen von Ressourcen für die Warteschlange direkt einzubinden. Verwenden Sie multihost_job.py als Arbeitsbeispiel; es löst den API-Aufruf zum Erstellen von Ressourcen für die Warteschlange mit dem entsprechenden Startscript aus, um Ihr Training auszuführen und nach einer vorzeitigen Beendigung fortzusetzen. Das Script multihost_job.py ist in der README-Datei für MaxText dokumentiert.

Da für multihost_job.py bei jedem Lauf 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 lieber direkt mit Befehlen für Ressourcen in der Warteschlange bereitstellen und verwalten möchten, lesen Sie Multislice-Umgebung verwalten.

Multislice-Umgebung verwalten

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

Ressourcen für die Warteschlange erstellen

gcloud

1. Erstellen Sie die Anfrage für in die Warteschlange gestellte Ressourcen mit dem folgenden Befehl.

   $ gcloud compute tpus queued-resources \
       create ${QR_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|--spot]

Achten Sie darauf, dass Sie das entsprechende Kontingent haben, bevor Sie --reserved, --spot oder das standardmäßige On-Demand-Kontingent auswählen. Informationen zu Kontingenttypen finden Sie auf der Seite zu den TPU-Kontingenten.

curl

1. Erstellen Sie eine Datei namens queued-resource-req.json und kopieren Sie den folgenden JSON-Code 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: Nummer des Projekts in Google Cloud
   • your-zone: Zone, in der Sie die in die Warteschlange gestellte Ressource erstellen möchten
   • accelerator-type: Version und Größe eines einzelnen Slice. Multislice wird nur auf Cloud TPU v4 und späteren TPU-Versionen unterstützt.
   • tpu-vm-runtime-version: TPU-VM-Laufzeitversion, die Sie verwenden möchten
   • your-network-name (optional): Netzwerk, an das die in die Warteschlange gestellte Ressource angehängt wird
   • your-subnetwork-name (optional): Subnetzwerk, an das die in die Warteschlange gestellte Ressource angehängt wird
   • example-tag-1 (optional): Beliebiger Tag-String
   • your-startup-script: Startscript, das ausgeführt wird, wenn die Ressource aus der Warteschlange zugewiesen wird
   • slice-count: Anzahl der TPU-Slices in Ihrer Multislice-Umgebung
   • your-queued-resource-id: Vom Nutzer bereitgestellte ID für die in die Warteschlange gestellte Ressource

   Weitere Informationen zu allen verfügbaren Optionen finden Sie in der Dokumentation zur REST API für „QueuedResource“.

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

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

   Entfernen Sie die Zeile, um die standardmäßige On-Demand-Kapazität zu verwenden.

2. Senden Sie die Anfrage für in die Warteschlange gestellte Ressourcen 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/your-project-id/locations/your-zone/queuedResources\?queued_resource_id\=your-queued-resource-id

   Ersetzen Sie die folgenden Werte:

   • your-project-id: Projekt-ID in Google Cloud
   • your-zone: Zone, in der Sie die in die Warteschlange gestellte Ressource erstellen möchten
   • your-queued-resource-id: Vom Nutzer bereitgestellte ID für die in die Warteschlange gestellte 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 Attribut name, um Informationen zur Anfrage für in die Warteschlange gestellte Ressourcen zu erhalten.

Console

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

   Zu TPUs

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 Sie die TPU erstellen möchten.

5. Wählen Sie im Feld TPU-Typ einen Beschleunigertyp aus. Der Beschleunigertyp gibt die Version und Größe der Cloud TPU an, die Sie erstellen möchten. Multislice wird nur auf Cloud TPU v4 und späteren TPU-Versionen unterstützt. Weitere Informationen zu den unterstützten Beschleunigertypen für die einzelnen TPU-Versionen finden Sie unter TPU-Versionen.

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

7. Klicken Sie auf die Umschaltfläche Warteschlange aktivieren.

8. Geben Sie im Feld Name der Ressource in der Warteschlange einen Namen für die Anfrage für in die Warteschlange gestellte Ressourcen ein.

9. Aktivieren Sie das Kästchen Multislice-TPU erstellen.

10. Geben Sie im Feld Slice-Anzahl die Anzahl der Slices ein, die Sie erstellen möchten.

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

Status einer in die Warteschlange gestellten Ressource abrufen

gcloud

$ gcloud compute tpus queued-resources describe ${QR_ID} --zone=${ZONE}

Für eine in die Warteschlange gestellte Ressource im Status ACTIVE sieht die Ausgabe in etwa 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}

Für eine in die Warteschlange gestellte Ressource im Status ACTIVE sieht die Ausgabe in etwa so aus:

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

Console

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

   Zu TPUs

2. Klicken Sie auf den Tab In die Warteschlange gestellte Ressourcen.

3. Klicken Sie auf den Namen Ihrer Anfrage für in die Warteschlange gestellte Ressourcen.

Nachdem Ihre TPU bereitgestellt wurde, können Sie auch Details zur Anfrage für in die Warteschlange gestellte Ressourcen aufrufen. Öffnen Sie dazu die Seite TPUs, suchen Sie nach Ihrer TPU und klicken Sie auf den Namen der zugehörigen Ressourcenanfrage.

In seltenen Fällen kann es vorkommen, dass sich Ihre in die Warteschlange gestellte Ressource im Status FAILED befindet, während einige Slices den Status ACTIVE haben. Löschen Sie in diesem Fall die erstellten Ressourcen und versuchen Sie es einige Minuten später noch einmal oder wenden Sie sich an den Support vonGoogle Cloud .

SSH-Verbindung herstellen und Abhängigkeiten installieren

Unter JAX-Code auf TPU-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 ${QR_ID} \
        --zone=${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 der Ressource gesendet. Der Befehl wird in Batches mit je vier Befehlen gebündelt und gleichzeitig gesendet. Der nächste Batch von Befehlen wird gesendet, nachdem der aktuelle Batch ausgeführt wurde. 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 in die Warteschlange gestellte Ressourcen. Wenn die Anzahl der verwendeten Slices das Threading-Limit (auch Batching-Limit genannt) Ihres lokalen Computers überschreitet, kommt es zu einem Deadlock. Angenommen, das Batching-Limit auf Ihrem lokalen Computer ist 64. Wenn Sie versuchen, ein Trainingsscript auf mehr als 64 Slices auszuführen, z. B. für 100, werden die Slices mit dem SSH-Befehl in Batches aufgeteilt. Das Trainingsscript wird auf dem ersten Batch mit 64 Slices ausgeführt. Das System wartet, bis die Scriptausführung abgeschlossen ist, bevor das Script auf dem verbleibenden Batch mit 36 Slices ausgeführt wird. Der erste Batch mit 64 Slices kann jedoch erst abgeschlossen werden, wenn die verbleibenden 36 Slices das Script ausführen, was zu einem Deadlock führt.

Um dieses Szenario zu vermeiden, können Sie das Trainingsscript auf jeder VM im Hintergrund ausführen. Hängen Sie dazu ein Und-Zeichen (&) an den Scriptbefehl an, den Sie mit dem Flag --command angeben. Die Steuerung wird dann nach dem Start des Trainingsscripts auf dem ersten Slices-Batch sofort an den SSH-Befehl zurückgegeben. Mit dem SSH-Befehl kann anschließend das Trainingsscript auf dem verbleibenden Batch mit 36 Slices ausgeführt werden. Sie müssen Ihre stdout- und stderr-Streams entsprechend weiterleiten, wenn Sie die Befehle im Hintergrund ausführen. Um die Parallelität innerhalb derselben Ressource in der Warteschlange zu erhöhen, können Sie mit dem Parameter --node bestimmte Slices auswählen.

Netzwerkeinrichtung

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

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

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

Ressourcen in der Warteschlange auflisten

gcloud

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

$ gcloud compute tpus queued-resources list --zone=${ZONE}

Die Ausgabe sieht dann ungefähr so aus:

NAME        ZONE           NODE_COUNT  ACCELERATOR_TYPE  STATE
...
que-res-id  us-central1-a  4           v5litepod-16             ACTIVE
...

Console

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

   Zu TPUs

2. Klicken Sie auf den Tab In die Warteschlange gestellte 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 ${TPU_NAME} \
    --zone=${ZONE} \
    --worker=all \
    --command="command-to-run"

In die Warteschlange gestellte Ressourcen zurücksetzen

Mit der ResetQueuedResource API können Sie alle VMs in einer in die Warteschlange gestellten Ressource mit dem Status ACTIVE zurücksetzen. Durch das Zurücksetzen der VMs wird der Arbeitsspeicher der Maschine gelöscht und die VM in ihren Ausgangszustand zurückgesetzt. Alle lokal gespeicherten Daten bleiben erhalten und nach dem Zurücksetzen wird das Startscript aufgerufen. Die ResetQueuedResource API kann nützlich sein, wenn Sie alle TPUs neu starten möchten. Das kann beispielsweise der Fall sein, wenn der Trainingsprozess nicht mehr reagiert und das Zurücksetzen aller VMs einfacher ist als das Debuggen.

Alle VMs werden parallel zurückgesetzt. Ein ResetQueuedResource-Vorgang dauert ein bis zwei Minuten. Verwenden Sie den folgenden Befehl, um die API aufzurufen:

$ gcloud compute tpus queued-resources reset ${QR_ID} --zone=${ZONE}

Ressourcen in der Warteschlange löschen

Um Ressourcen am Ende des Trainings freizugeben, löschen Sie die Ressource in der Warteschlange. 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 ${QR_ID} --zone=${ZONE} --force [--async]

Console

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

   Zu TPUs

2. Klicken Sie auf den Tab In die Warteschlange gestellte Ressourcen.

3. Aktivieren Sie das Kästchen neben Ihrer Anfrage für in die Warteschlange gestellte Ressourcen.

4. Klicken Sie auf delete Löschen.

Automatische Wiederherstellung nach Fehlern

Im Falle einer Störung bietet Multislice eine automatische Reparatur des betroffenen Slice und ein anschließendes Zurücksetzen aller Slices. Der betroffene Slice wird durch einen neuen ersetzt und die verbleibenden, ansonsten fehlerfreien Slices werden zurückgesetzt. Wenn keine Kapazität für die Zuweisung eines Ersatz-Slice verfügbar ist, 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 Prüfpunkten sucht und diese lädt. Ihr Startscript wird automatisch jedes Mal ausgeführt, wenn ein Slice neu zugewiesen oder eine VM zurückgesetzt wird. Sie geben ein Startscript in der JSON-Nutzlast an, die Sie an die API zum Erstellen einer Anfrage für Ressourcen in der Warteschlange senden.

Das folgende Startscript (das in Ressourcen für die Warteschlange erstellen verwendet wird) ermöglicht während des MaxText-Trainings das automatische Wiederherstellen nach Fehlern und Fortsetzen des Trainings anhand von Prüfpunkten, die 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 -m MaxText.train 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 das ausprobieren.

Profilerstellung und Debugging

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

Training optimieren

In den folgenden Abschnitten wird beschrieben, wie Sie das Multislice-Training optimieren können.

Fragmentierung mit Multislice für maximale Leistung

Um in Multislice-Umgebungen maximale Leistung zu erzielen, ist zu bedenken, wie die Daten am besten auf die verschiedenen Slices aufgeteilt (fragmentiert) werden. Normalerweise gibt es drei Möglichkeiten: Datenparallelität, vollständig fragmentierte Datenparallelität und Pipelineparallelität. Die Fragmentierung von Aktivierungen über die Modelldimensionen hinweg (manchmal auch als Tensorparallelität bezeichnet) ist nicht zu empfehlen, da dies zu viel Bandbreite zwischen den Slices erfordert. Sie können für einen Slice immer diejenige dieser Fragmentierungsstrategien beibehalten, die in der Vergangenheit für Sie funktioniert hat.

Wir empfehlen, mit reiner Datenparallelität zu beginnen. Die vollständig fragmentierte Datenparallelität ist hilfreich, um die Arbeitsspeichernutzung zu reduzieren. Der Nachteil ist, dass für die Kommunikation zwischen den Slices das DCN-Netzwerk verwendet wird und diese Kommunikation Ihre Arbeitslast verlangsamt. Verwenden Sie Pipelineparallelität nur, wenn es aufgrund der Batchgröße erforderlich ist (siehe Analyse unten).

Wann sollte die Datenparallelität verwendet werden?

Reine Datenparallelität eignet sich gut für Arbeitslasten, die bereits zufriedenstellend laufen, deren Leistung Sie aber durch Skalierung über mehrere Slices verbessern möchten.

Um eine starke Skalierung über mehrere Slices zu erreichen, muss ein All-Reduce-Vorgang über das DCN schneller möglich sein als ein Rückwärtsdurchlauf. Das DCN wird für die Kommunikation zwischen Slices verwendet und ist ein einschränkender Faktor für den Arbeitslastdurchsatz.

Jeder v4-TPU-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 pro Schritt 32 bis 64 Bit DCN-Bandbreite für jeden Parameter. Wenn Sie zwei Slices verwenden, nutzt Ihr Modell 32 Bit DCN-Bandbreite. Wenn Sie mehr als zwei Slices verwenden, führt der Compiler einen Full-Shuffle-All-Reduce-Vorgang aus und Sie verbrauchen pro Schritt bis zu 64 Bit DCN-Bandbreite für jeden Parameter. Die Anzahl der FLOPS, die für jeden Parameter erforderlich sind, hängt vom jeweiligen Modell ab. Insbesondere für Transformer-basierte Sprachmodelle ist die Anzahl der FLOPS, die für einen Vorwärts- und einen Rückwärtsdurchlauf erforderlich sind, ungefähr 6 * B * P, wobei gilt:

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

Um eine starke Skalierung über mehrere Slices zu erreichen, muss die operative Intensität die arithmetische Intensität der TPU-Hardware übersteigen. Um die operative Intensität zu berechnen, 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 also für ein Transformer-basiertes Sprachmodell zwei Slices verwenden: Operational intensity = 4 * B / 32

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

Daraus ergibt sich eine empfehlenswerte Mindestbatchgröße zwischen 176.000 und 352.000 für Transformer-basierte Sprachmodelle. Da im DCN Pakete kurzzeitig verloren gehen können, ist es am besten, eine große Fehlertoleranz zu verwenden und Datenparallelität nur einzusetzen, 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 Slice schätzen (entweder durch Messung mit einem Profiler oder durch Zählen von FLOPS). Anschließend können Sie diese mit der erwarteten Laufzeit für den All-Reduce-Vorgang über das DCN vergleichen. Dann können Sie besser einschätzen, ob sich Datenparallelität für Sie lohnt.

Wann sollte FSDP (Fully Sharded Data Parallelism) verwendet werden?

Bei der vollständig fragmentierten Datenparallelität (Fully Sharded Data Parallelism, FSDP) wird die Datenparallelität (Fragmentierung der Daten über mehrere Knoten) mit der Fragmentierung der Gewichtungen über mehrere Knoten kombiniert. Die Gewichtungen für jeden Vorgang in Vorwärts- und Rückwärtsdurchläufen werden mit „all-gather“ erfasst, sodass jeder Slice die benötigten Gewichtungen hat. Anstatt die Gradienten mit „all-reduce“ zu synchronisieren, werden sie mit „reduce-scatter“ reduziert und verteilt, sobald sie erzeugt werden. So erhält jeder Slice nur die Gradienten für die jeweils relevanten Gewichtungen.

Ähnlich wie bei der Datenparallelität muss die globale Batchgröße bei FSDP linear mit der Anzahl der Slices skaliert werden. Durch FSDP wird die Arbeitsspeicherlast verringert, wenn Sie die Anzahl der Slices erhöhen. Das liegt daran, dass die Anzahl der Gewichtungen und der Optimiererstatus pro Slice abnimmt, dies jedoch auf Kosten von erhöhtem Netzwerkverkehr und einer größeren Gefahr von Blockierungen aufgrund eines verzögerten kollektiven Vorgangs geschieht.

In der Praxis funktioniert FSDP über mehrere Slices am besten, wenn Sie die Batchgröße pro Slice erhöhen, mehr Aktivierungen speichern, um die Rematerialisierung während des Rückwärtsdurchlaufs zu minimieren, oder die Anzahl der Parameter im neuronalen Netzwerk erhöhen.

Die Vorgänge „all-gather“ und „all-reduce“ in FSDP funktionieren ähnlich wie bei DP. Sie können also auf dieselbe Weise wie im vorherigen Abschnitt beschrieben feststellen, ob Ihre FSDP-Arbeitslast durch die DCN-Leistung eingeschränkt wird.

Wann sollte die Pipelineparallelität verwendet werden?

Die Pipelineparallelität wird relevant, wenn Sie mit anderen Parallelitätsstrategien nur eine hohe Leistung erzielen können, wenn die globale Batchgröße über der von Ihnen bevorzugten maximalen Batchgröße liegt. Mit der Pipelineparallelität können die Slices einer Pipeline einen Batch gemeinsam nutzen. Die Pipelineparallelität hat jedoch zwei erhebliche Nachteile:

1. Es kommt zur „Pipeline-Blase“, bei der Chips inaktiv sind, weil sie auf Daten warten.
2. Außerdem ist Micro-Batching erforderlich, wodurch die effektive Batchgröße, die arithmetische Intensität und letztendlich die FLOP-Auslastung des Modells verringert werden.

Die Pipelineparallelität 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, die für eine leistungsstarke FSDP erforderlich ist, verlangsamt. Mit FSDP wird in der Regel eine höhere FLOP-Auslastung des Modells erreicht. Wenn sich die Konvergenz pro Stichprobe jedoch mit zunehmender Batchgröße verlangsamt, ist Pipelineparallelität möglicherweise trotzdem die bessere Wahl. Die meisten Arbeitslasten können so große Batchgrößen tolerieren, dass sie nicht von Pipelineparallelität profitieren. Bei Ihrer Arbeitslast kann es jedoch anders sein.

Wenn Pipelineparallelität erforderlich ist, empfehlen wir, sie mit Datenparallelität oder FSDP zu kombinieren. So können Sie die Tiefe der Pipeline minimieren und gleichzeitig die Batchgröße pro Pipeline erhöhen, bis die DCN-Latenz weniger Einfluss auf den Durchsatz hat. Konkret: Bei N Slices empfehlen sich Pipelines mit Tiefe 2 und N/2 Replikaten für Datenparallelität, dann Pipelines mit Tiefe 4 und N/4 Replikaten für Datenparallelität usw., bis der Batch pro Pipeline so groß wird, dass die kollektiven DCN-Vorgänge hinter der Arithmetik im Rückwärtsdurchlauf verborgen werden können. Dadurch wird die durch Pipelineparallelität verursachte Verlangsamung minimiert und Sie können über das globale Batchgrößenlimit hinaus skalieren.

Best Practices für Multislice

In den folgenden Abschnitten werden Best Practices für das Multislice-Training beschrieben.

Laden der Daten

Während des Trainings werden wiederholt Batches aus einem Dataset geladen, um sie in das Modell einzuspeisen. Ein effizienter, asynchroner Data Loader, der den Batch auf Hosts verteilt, ist wichtig, um unterausgelastete TPUs zu vermeiden. Der aktuelle Data Loader in MaxText lässt jeden Host eine gleiche Teilmenge der Beispiele laden. Diese Lösung ist für Text geeignet, erfordert jedoch eine erneute Fragmentierung im Modell. Außerdem bietet MaxText noch keine deterministischen Snapshots. Mit einer solchen Funktion könnte der Dateniterator vor und nach einer vorzeitigen Beendigung dieselben Daten laden.

Prüfpunkte

Die Orbax-Bibliothek bietet Bausteine zur Erstellung von Prüfpunkten von JAX-PyTrees in lokalen Speichern oder Speicher in Google Cloud . Wir bieten eine Referenzintegration mit synchroner Prüfpunkterstellung in MaxText in checkpointing.py.

Unterstützte Konfigurationen

In den folgenden Abschnitten werden die unterstützten Slice-Formen, Orchestrierung, Frameworks und Parallelität für Multislice beschrieben.

Formen

Alle Slices müssen dieselbe Form haben (z. B. denselben AcceleratorType). Heterogene Slice-Formen werden nicht unterstützt.

Orchestrierung

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

Frameworks

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

Parallelität

Wir empfehlen, Multislice mit Datenparallelität zu testen. Weitere Informationen zur Implementierung von Pipelineparallelität mit Multislice erhalten Sie von Ihrem Kundenbetreuer beiGoogle Cloud .

Support und Feedback

Wir freuen uns über jedes Feedback. Wenn Sie Feedback geben oder Support anfordern möchten, wenden Sie sich über das Support- oder Feedbackformular für Cloud TPU an uns.