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.

Mehrere Slices skalieren die Leistung linear.

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.

Multislice-Datenfluss

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.

    3. 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