TPU v4

In diesem Dokument werden die Architektur und die unterstützten Konfigurationen von Cloud TPU v4 beschrieben.

Systemarchitektur

Jeder TPU v4-Chip enthält zwei TensorCores. Jeder TensorCore hat vier Matrixmultiplikationseinheiten (MXUs), eine Vektoreinheit und eine Skalareinheit. In der folgenden Tabelle sind die wichtigsten Spezifikationen für einen TPU-Pod der Version 4 aufgeführt.

Das folgende Diagramm zeigt einen TPU v4-Chip.

Weitere Informationen zu Architekturdetails und Leistungsmerkmalen von TPU v4 finden Sie unter TPU v4: Ein optisch rekonfigurierbarer Supercomputer für maschinelles Lernen mit Hardwareunterstützung für Einbettungen.

3D-Mesh und 3D-Torus

TPUs der Version 4 haben eine direkte Verbindung zu den nächstgelegenen benachbarten Chips in drei Dimensionen, was zu einem 3D-Mesh aus Netzwerkverbindungen führt. Die Verbindungen können als 3D-Torus in Slices konfiguriert werden, bei denen die Topologie AxBxC entweder 2A=B=C oder 2A=2B=C ist, wobei jede Dimension ein Vielfaches von 4 ist. Beispiel: 4 × 4 × 8, 4 × 8 × 8 oder 12 × 12 × 24. Im Allgemeinen ist die Leistung einer 3D-Toruskonfiguration besser als die einer 3D-Mesh-Konfiguration. Weitere Informationen finden Sie unter Twisted-Torus-Topologien.

Leistungsvorteile von TPU v4 im Vergleich zu v3

In diesem Abschnitt wird eine speichereffiziente Methode zum Ausführen eines Beispiel-Trainingsscripts auf TPU v4 beschrieben. Außerdem werden die Leistungsverbesserungen von TPU v4 im Vergleich zu TPU v3 gezeigt.

Speichersystem

Non Uniform Memory Access (NUMA) ist eine Computerspeicherarchitektur für Maschinen mit mehreren CPUs. Jede CPU hat direkten Zugriff auf einen Block mit Hochgeschwindigkeitsspeicher. Eine CPU und ihr Arbeitsspeicher werden als NUMA-Knoten bezeichnet. NUMA-Knoten sind mit anderen NUMA-Knoten verbunden, die direkt nebeneinander liegen. Eine CPU eines NUMA-Knotens kann auf den Arbeitsspeicher eines anderen NUMA-Knotens zugreifen. Dieser Zugriff ist jedoch langsamer als der Zugriff auf den Arbeitsspeicher innerhalb eines NUMA-Knotens.

Software, die auf einem Computer mit mehreren CPUs ausgeführt wird, kann Daten, die von einer CPU benötigt werden, in ihrem NUMA-Knoten platzieren und so den Arbeitsspeicherdurchsatz erhöhen. Weitere Informationen zu NUMA finden Sie auf Wikipedia unter Non Uniform Memory Access.

Sie können die Vorteile der NUMA-Lokalität nutzen, indem Sie Ihr Trainingsskript an NUMA-Knoten 0 binden.

So aktivieren Sie die NUMA-Knotenbindung:

1. Erstellen Sie eine TPU v4.

2. Stellen Sie eine SSH-Verbindung zur TPU her.

3. Installieren Sie das Befehlszeilentool numactl. Mit numactl können Sie Prozesse mit einer bestimmten NUMA-Planungs- oder Speicherplatzierungsrichtlinie ausführen.

    $ sudo apt-get update
    $ sudo apt-get install numactl

4. Binden Sie Ihren Scriptcode an NUMA-Knoten 0. Ersetzen Sie your-training-script durch den Pfad zu Ihrem Trainingsskript.

    $ numactl --cpunodebind=0 python3 your-training-script

Aktivieren Sie die NUMA-Knotenbindung, wenn:

• Wenn Ihre Arbeitslast unabhängig vom Framework stark von CPU-Arbeitslasten abhängt (z. B. Bildklassifizierung, Empfehlungsarbeitslasten).
• Wenn Sie eine TPU-Laufzeitversion ohne Suffix „-pod“ verwenden (z. B. tpu-vm-tf-2.10.0-v4).

Weitere Unterschiede beim Arbeitsspeichersystem:

• TPU-v4-Chips haben einen einheitlichen HBM-Speicherplatz von 32 GiB auf dem gesamten Chip, was eine bessere Koordination zwischen den beiden On-Chip-TensorCores ermöglicht.
• Verbesserte HBM-Leistung durch die neuesten Speicherstandards und -geschwindigkeiten.
• Verbessertes DMA-Leistungsprofil mit integrierter Unterstützung für hochleistungsfähiges Striding mit einer Granularität von 512 B.

Tensorkerne

• Doppelt so viele MXUs und eine höhere Taktrate sorgen für 275 TFLOPS.
• 2-fache Transposition- und Permutationsbandbreite.
• Load-Store-Speicherzugriffsmodell für Common Memory (Cmem).
• Höhere Bandbreite für das Laden von MXU-Gewichten und Unterstützung des 8‑Bit-Modus für kleinere Batch-Größen und eine verbesserte Inferenzlatenz.

Chip-zu-Chip-Verbindung

Sechs Interconnect-Glieder pro Chip für Netzwerktopologien mit kleinerem Netzwerkdurchmesser

Sonstiges

• x16-PCIe-Gen3-Schnittstelle zum Host (direkte Verbindung).
• Verbessertes Sicherheitsmodell.
• Verbesserte Energieeffizienz.

Konfigurationen

Ein TPU v4-Pod besteht aus 4.096 Chips, die über rekonfigurierbare Hochgeschwindigkeits-Glieder miteinander verbunden sind. Dank des flexiblen Netzwerks von TPU v4 können Sie die Chips in einem Slice derselben Größe auf verschiedene Arten verbinden. Beim Erstellen eines TPU-Slice geben Sie die TPU-Version und die Anzahl der erforderlichen TPU-Ressourcen an. Wenn Sie ein TPU v4-Speichersegment erstellen, können Sie dessen Typ und Größe auf zwei Arten angeben: AcceleratorType und AccleratorConfig.

AcceleratorType verwenden

Verwenden Sie AcceleratorType, wenn Sie keine Topologie angeben. Wenn Sie TPUs der Version 4 mit AcceleratorType konfigurieren möchten, verwenden Sie beim Erstellen des TPU-Slabs das Flag --accelerator-type. Legen Sie --accelerator-type auf einen String fest, der die TPU-Version und die Anzahl der zu verwendenden TensorCores enthält. Wenn Sie beispielsweise ein v4-Speicher-Slice mit 32 TensorCores erstellen möchten, verwenden Sie --accelerator-type=v4-32.

Verwenden Sie den Befehl gcloud compute tpus tpu-vm create mit dem Flag --accelerator-type, um ein TPU-Slice der Version 4 mit 512 TensorCores zu erstellen:

  $ gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-central2-b \
    --accelerator-type=v4-512 \
    --version=tpu-ubuntu2204-base

Die Zahl nach der TPU-Version (v4) gibt die Anzahl der Tensorcores an. Eine v4-TPU hat zwei TensorCores, sodass die Anzahl der TPU-Chips 512 ÷ 2 = 256 beträgt.

Weitere Informationen zum Verwalten von TPUs finden Sie unter TPUs verwalten. Weitere Informationen zur Systemarchitektur von Cloud TPU finden Sie unter Systemarchitektur.

AcceleratorConfig verwenden

Verwenden Sie AcceleratorConfig, wenn Sie die physische Topologie Ihres TPU-Slabs anpassen möchten. Dies ist in der Regel für die Leistungsoptimierung bei Slices mit mehr als 256 Chips erforderlich.

Wenn Sie TPUs der Version 4 mit AcceleratorConfig konfigurieren möchten, verwenden Sie die Flags --type und --topology. Legen Sie --type auf die gewünschte TPU-Version und --topology auf die physische Anordnung der TPU-Chips im Slice fest.

Sie geben eine TPU-Topologie mit einem Dreier-Tupel an, AxBxC, wobei A<=B<=C und A, B, C entweder alle <= 4 oder alle ganzzahligen Vielfache von 4 sind. Die Werte A, B und C sind die Chipanzahl in den drei Dimensionen. Wenn Sie beispielsweise ein v4-Speicher-Slice mit 16 Chips erstellen möchten, legen Sie --type=v4 und --topology=2x2x4 fest.

Verwenden Sie den Befehl gcloud compute tpus tpu-vm create, um ein TPU v4-Speicher-Slice mit 128 TPU-Chips in einem 4 × 4 × 8-Array zu erstellen:

  $ gcloud compute tpus tpu-vm create your-tpu-name \
    --zone=us-central2-b \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-ubuntu2204-base

Für Topologien mit 2 A=B=C oder 2 A=2 B=C gibt es auch Topologievarianten, die für die All-to-All-Kommunikation optimiert sind, z. B. 4 × 4 × 8, 8 × 8 × 16 und 12 × 12 × 24. Diese werden als verdrillte Torus-Topologien bezeichnet.

Die folgenden Abbildungen zeigen einige gängige TPU v4-Topologien.

Größere Scheiben können aus einem oder mehreren 4 × 4 × 4 Chip-Würfeln aufgebaut werden.

Weitere Informationen zum Verwalten von TPUs finden Sie unter TPUs verwalten. Weitere Informationen zur Systemarchitektur von Cloud TPU finden Sie unter Systemarchitektur.

Twisted-Torus-Topologien

Bei einigen 3D-Torus-Scheibenformen der Version 4 kann die sogenannte verdrehte Torus-Topologie verwendet werden. Zwei V4-Würfel können beispielsweise als 4x4x8-Segment oder 4x4x8_twisted angeordnet werden. Verdrillte Topologien bieten eine deutlich höhere Bisection-Bandbreite. Ein Slice mit der Topologie „4x4x8_twisted“ bietet beispielsweise eine theoretische Steigerung der Bisection-Bandbreite um 70 % im Vergleich zu einem nicht verdrillten 4x4x8-Slice. Eine erhöhte Bisection-Bandbreite ist für Arbeitslasten nützlich, bei denen globale Kommunikationsmuster verwendet werden. Verdrillte Topologien können die Leistung der meisten Modelle verbessern. Am meisten profitieren große TPU-Embedding-Arbeitslasten.

Bei Arbeitslasten, bei denen die Datenparallelität die einzige Parallelitätsstrategie ist, können verdrillte Topologien eine etwas bessere Leistung erzielen. Bei LLMs kann die Leistung bei Verwendung einer verdrillten Topologie je nach Art der Parallelität (DP, MP usw.) variieren. Am besten trainieren Sie Ihr LLM mit und ohne verdrillte Topologie, um festzustellen, welche die beste Leistung für Ihr Modell bietet. Bei einigen Tests mit dem FSDP MaxText-Modell wurden mit einer verdrillten Topologie Verbesserungen von 1–2 MFU erzielt.

Der Hauptvorteil von verdrillten Topologien besteht darin, dass eine asymmetrische Torus-Topologie (z. B. 4 × 4 × 8) in eine eng verwandte symmetrische Topologie umgewandelt wird. Die symmetrische Topologie bietet viele Vorteile:

• Verbessertes Load Balancing
• Höhere Bisection-Bandbreite
• Kürzere Paketpfade

Diese Vorteile führen letztendlich zu einer verbesserten Leistung bei vielen globalen Kommunikationsmustern.

Die TPU-Software unterstützt verdrillte Torus-Objekte in Scheiben, bei denen die Größe jeder Dimension entweder der kleinsten Dimension entspricht oder doppelt so groß ist. Beispiel: 4 × 4 × 8, 4 × 8 × 8 oder 12 × 12 × 24.

Betrachten Sie als Beispiel diese 4×2-Torus-Topologie mit TPUs, die im Slice mit ihren (X,Y)-Koordinaten gekennzeichnet sind:

Die Kanten in diesem Topologiediagramm werden zur Verdeutlichung als ungerichtete Kanten dargestellt. In der Praxis ist jeder Edge eine bidirektionale Verbindung zwischen TPUs. Die Ränder zwischen einer Seite dieses Rasters und der gegenüberliegenden Seite werden als Umlaufränder bezeichnet, wie im Diagramm dargestellt.

Durch das Verdrehen dieser Topologie erhalten wir eine vollständig symmetrische 4×2-verdrehte Torus-Topologie:

Der einzige Unterschied zwischen diesem Diagramm und dem vorherigen sind die Y-Achsenränder. Anstatt eine Verbindung zu einer anderen TPU mit derselben X-Koordinate herzustellen, wurden sie verschoben, um eine Verbindung zur TPU mit der Koordinate X + 2 mod 4 herzustellen.

Das Gleiche gilt für unterschiedliche Dimensionengrößen und unterschiedliche Anzahlen von Dimensionen. Das resultierende Netzwerk ist symmetrisch, solange jede Dimension der kleinsten Dimension entspricht oder doppelt so groß ist.

Unter AcceleratorConfig verwenden finden Sie weitere Informationen dazu, wie Sie beim Erstellen einer Cloud TPU eine verdrillte Torus-Konfiguration angeben.

In der folgenden Tabelle sind die unterstützten verdrillten Topologien und eine theoretische Steigerung der Bisection-Bandbreite im Vergleich zu unverdrillten Topologien aufgeführt.

TPU v4-Topologievarianten

Einige Topologien mit derselben Anzahl von Chips können auf unterschiedliche Weise angeordnet werden. Ein TPU-Segment mit 512 Chips (1.024 TensorCores) kann beispielsweise mit den folgenden Topologien konfiguriert werden: 4 x 4 x 32, 4 x 8 x 16 oder 8 x 8 x 8. Ein TPU-Segment mit 2.048 Chips (4.096 TensorCores) bietet noch mehr Topologieoptionen: 4x4x128, 4x8x64, 4x16x32 und 8x16x16.

Die Standardtopologie, die mit einer bestimmten Chipanzahl verknüpft ist, ist diejenige, die einem Würfel am ähnlichsten ist. Diese Form ist wahrscheinlich die beste Wahl für das datenparrallele ML-Training. Andere Topologien können für Arbeitslasten mit mehreren Arten von Parallelität nützlich sein, z. B. Modell- und Datenparallelität oder räumliche Partitionierung einer Simulation. Diese Arbeitslasten erzielen die beste Leistung, wenn die Topologie an die verwendete Parallelität angepasst ist. Wenn Sie beispielsweise eine 4‑Wege-Modellparallelität auf die X‑Achse und eine 256‑Wege-Datenparallelität auf die Y‑ und Z‑Achse anwenden, entspricht das einer 4x16x16-Topologie.

Bei Modellen mit mehreren Parallelitätsdimensionen ist die Leistung am besten, wenn die Parallelitätsdimensionen den TPU-Topologiedimensionen zugeordnet sind. Das sind in der Regel parallele Daten- und Modell-Large Language Models (LLMs). Bei einem TPU v4-Segment mit der Topologie 8x16x16 sind die TPU-Topologieabmessungen beispielsweise 8, 16 und 16. Die Leistung ist höher, wenn Sie eine 8‑ oder 16‑fach-Modellparallelität verwenden, die einer der physischen TPU-Topologiedimensionen zugeordnet ist. Eine 4-Wege-Modellparallelität wäre mit dieser Topologie nicht optimal, da sie nicht mit einer der TPU-Topologiedimensionen übereinstimmt. Sie wäre jedoch optimal mit einer 4x16x32-Topologie bei derselben Anzahl von Chips.

TPU v4-Konfigurationen bestehen aus zwei Gruppen: Topologien mit weniger als 64 Chips (kleine Topologien) und Topologien mit mehr als 64 Chips (große Topologien).

Kleine v4-Topologien

Cloud TPU unterstützt die folgenden TPU v4-Segmente mit weniger als 64 Chips, also einen Würfel mit 4 × 4 × 4 Chips. Sie können diese kleinen v4-Topologien entweder anhand ihres TensorCore-basierten Namens (z. B. v4-32) oder ihrer Topologie (z. B. 2x2x4) erstellen:

Große v4-Topologien

TPU v4-Segmente sind in Schritten von 64 Chips verfügbar. Die Form ist in allen drei Dimensionen ein Vielfaches von 4. Die Abmessungen müssen in aufsteigender Reihenfolge vorliegen. In der folgenden Tabelle finden Sie einige Beispiele. Einige dieser Topologien sind „benutzerdefinierte“ Topologien, die nur mit den Flags --type und --topology erstellt werden können, da es mehr als eine Möglichkeit gibt, die Chips anzuordnen.