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.

Wichtige Spezifikationen	v4-Pod-Werte
Spitzenleistung pro Chip	275 Teraflops (bf16 oder int8)
HBM2-Kapazität und ‑Bandbreite	32 GiB, 1.200 GBps
Gemessene minimale/durchschnittliche/maximale Leistung	90/170/192 W
TPU-Pod-Größe	4.096 Chips
Interconnect-Topologie	3D-Mesh
Spitzenberechnung pro Pod	1,1 Exaflops (bf16 oder int8)
All-Reduce-Bandbreite pro Pod	1,1 PB/s
Geteilte Bandbreite pro Pod	24 TB/s

Das folgende Diagramm zeigt einen TPU v4-Chip.

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 werden die Leistungsvorteile von TPU v4 beschrieben.

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. Installieren Sie das Befehlszeilentool „numactl“.

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

2. Binden Sie Ihren Scriptcode an NUMA-Knoten 0.

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

TensorCores

• 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 Batchgrößen und eine verbesserte Inferenzlatenz.

Chip-zu-Chip-Verbindung

Sechs Interconnect-Glieder pro Chip ermöglichen Netzwerktopologien mit kleinerem Netzwerkdurchmesser.

Sonstiges

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

Konfigurationen

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

Hinweis: AcceleratorType und AccleratorConfig schließen sich gegenseitig aus. Sie können beim Erstellen einer TPU nur eines der beiden Flags verwenden.

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-Pod-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-Pod-Slice mit 32 TensorCores erstellen möchten, verwenden Sie --accelerator-type=v4-32.

Mit dem folgenden Befehl wird ein TPU-Pod-Slice der Version 4 mit 512 TensorCores mit dem Flag --accelerator-type erstellt:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --accelerator-type=v4-512 \
    --version=tpu-vm-tf-2.18.0-pod-pjrt

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

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 Pod-Slices mit mehr als 256 Chips erforderlich.

Wenn Sie TPUs der Version 4 mit AcceleratorConfig konfigurieren möchten, verwenden Sie die Flags --version und --topology. Legen Sie --version auf die gewünschte TPU-Version und --topology auf die physische Anordnung der TPU-Chips im Pod-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 einen v4-Pod-Speicherbereich mit 16 Chips erstellen möchten, setzen Sie --version=v4 und --topology=2x2x4.

Der folgende Befehl erstellt ein TPU-Pod-Stück der Version 4 mit 128 TPU-Chips, die in einem 4 × 4 × 8-Array angeordnet sind:

  $ gcloud compute tpus tpu-vm create tpu-name \
    --zone=zone \
    --type=v4 \
    --topology=4x4x8 \
    --version=tpu-vm-tf-2.18.0-pod-pjrt

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 Pod-Slices können aus einem oder mehreren 4 × 4 × 4 Chip-Würfeln aufgebaut werden.

Twisted-Torus-Topologien

Bei einigen 3D-Torus-Scheibenformen der Version 4 kann die sogenannte verdrehte Torus-Topologie verwendet werden. Beispielsweise können zwei V4-Würfel als 4x4x8-Segment oder 4x4x8_twisted angeordnet werden. Verdrillte Topologien bieten eine deutlich höhere Bisection-Bandbreite. 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 beispielsweise 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 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.

Twisted-Topologie	Theoretische Erhöhung der Bandbreite bei Bisection im Vergleich zu einem nicht verdrillten Torus
4×4×8_twisted	~70%
8x8x16_twisted
12×12×24_twisted
4×8×8_twisted	~40%
8×16×16_twisted

Hinweis: Die verdrillte Topologie wird nicht automatisch verwendet, es sei denn, das Suffix _twisted wird mit dem Flag --topology angegeben. Wenn Sie beispielsweise eine verdrillte 4x4x8-Topologie angeben möchten, geben Sie das Flag topology als 4x4x8_twisted an.

TPU v4-Topologievarianten

Einige Topologien mit derselben Anzahl von Chips können auf unterschiedliche Weise angeordnet werden. Ein TPU-Pod-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-Pod-Slice 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-Pod-Speicherbereich 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-Pod-Segmente mit weniger als 64 Chips, also einem 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:

Name (basierend auf der Anzahl der TensorCores)	Anzahl der Chips	Topologie
v4-8	4	2x2x1
v4-16	8	2x2x2
v4-32	16	2x2x4
v4-64	32	2x4x4

Große v4-Topologien

TPU v4-Pod-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 angegeben werden. 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.

Name (basierend auf der Anzahl der TensorCores)	Anzahl der Chips	Topologie
v4-128	64	4x4x4
v4-256	128	4x4x8
v4-512	256	4x8x8
Benutzerdefinierte Topologie: Die Flags --type und --topology müssen verwendet werden.	256	4x4x16
v4-1024	512	8 × 8 × 8
v4-1536	768	8 x 8 x 12
v4-2048	1.024	8x8x16
Benutzerdefinierte Topologie: Es müssen die Flags --type und --topology verwendet werden.	1.024	4x16x16
v4-4096	2.048	8x16x16
…	…	…