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 einer Skalareinheit. In der folgenden Tabelle sind die wichtigsten Spezifikationen für einen TPU-Pod der Version 4 aufgeführt.

Wichtige Spezifikationen	v4-Pod-Werte
Spitzenwert der Rechenleistung pro Chip	275 TeraFLOPS (bf16 oder int8)
HBM2-Kapazität und -Bandbreite	32 GiB, 1.200 GB/s
Gemessene minimale/mittlere/maximale Leistung	90/170/192 W
TPU-Pod-Größe	4096-Chips
Interconnect-Topologie	3D-Mesh
Spitzenwert für die Berechnung pro Pod	1.1 ExaFLOPS (bf16 oder int8)
All-Reduce-Bandbreite pro Pod	1,1 PB/s
Biabschnitte 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 x 4 x 8, 4 x 8 x 8 oder 12 x 12 x 24. Im Allgemeinen ist die Leistung eines 3D-Torus ist besser als eine 3D-Mesh-Konfiguration. Weitere Informationen Siehe Twisted-tori-Topologien.

Leistungsvorteile von TPU v4 gegenüber 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 aus einem NUMA-Knoten kann auf den Arbeitsspeicher eines anderen NUMA-Knotens zugreifen, aber dieser Zugriff ist langsamer als der Zugriff in einem NUMA-Knoten.

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 Siehe Non Uniform Memory Access. auf Wikipedia.

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 numactl-Befehlszeilentool.

    $ 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 im Speichersystem:

• 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 hohe Leistung mit 512 Mrd. Detaillierungsgraden.

Tensorkerne

• Doppelt so viele MXUs und eine höhere Taktrate sorgen für 275 TFLOPS.
• Doppelte Bandbreite für Transposition und Permutation.
• Load-Store-Arbeitsspeicherzugriffsmodell 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.

Inter-Chip-Interconnect

Sechs Interconnect-Verbindungen pro Chip für Netzwerktopologien mit kleineren Netzwerkdurchmesser.

Sonstiges

• x16-PCIE-gen3-Schnittstelle zum Host (Direct Connect).
• Verbessertes Sicherheitsmodell.
• Höhere Energieeffizienz.

Konfigurationen

Ein TPU v4 Pod besteht aus 4.096 Chips, die miteinander verbunden sind. Hochgeschwindigkeits-Links. Dank des flexiblen Netzwerks von TPU v4 können Sie die Chips gleich großen Pod-Slice auf mehrere Arten. Wenn Sie ein TPU-Pod-Slice erstellen, 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 verwenden.

AcceleratorType verwenden

Verwenden Sie AcceleratorType, wenn Sie keine Topologie angeben. v4-TPUs konfigurieren Verwenden Sie AcceleratorType und das Flag --accelerator-type beim Erstellen der TPU-Pod-Slice. 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.17.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 wäre 512/2 = 256.

AcceleratorConfig verwenden

Verwenden Sie AcceleratorConfig, wenn Sie die physische Topologie anpassen möchten Ihres TPU-Slice. Dies ist im Allgemeinen für die Leistungsoptimierung mit Pod erforderlich mit mehr als 256 Chips.

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

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. Um beispielsweise einen v4-Pod zu erstellen, mit 16 Chips, würden Sie --version=v4 und --topology=2x2x4 festlegen.

Mit dem folgenden Befehl wird ein v4-TPU-Pod-Slice mit 128 TPU-Chips erstellt, ein 4x4x8-Array:

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

Topologien, bei denen 2A=B=C oder 2A=2B=C ist, haben auch Topologievarianten. optimiert für die gesamte Kommunikation, 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-Slice-Formen in Version 4 kann eine sogenannte Twisted Torus-Topologie. Beispielsweise können zwei V4-Würfel als 4x4x8-Segment oder 4x4x8_twisted angeordnet werden. Verdrehte Topologien bieten eine deutlich höhere Teilung Bandbreite. Eine höhere geteilte Bandbreite ist nützlich für Arbeitslasten, die globale Kommunikationsmuster. Verdrehte Topologien verbessern bei den meisten Modelle mit großen TPU-Einbettungsarbeitslasten am meisten davon profitiert.

Bei Arbeitslasten, bei denen die Datenparallelität die einzige Parallelitätsstrategie ist, können verdrillte Topologien eine etwas bessere Leistung erzielen. Für LLMs ist die Leistung mit einem Twisted Topologie kann je nach Art der Parallelität (DP, MP usw.) variieren. Beste üben, Ihr LLM mit und ohne verdrillte Topologie zu trainieren, die 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 hat viele Vorteile:

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

Diese Vorteile führen letztendlich zu einer verbesserten Leistung vieler globaler Kommunikationsmuster.

Die TPU-Software unterstützt verdrehte Tore auf Segmenten, bei denen die Größe der einzelnen Segmente Dimension entweder gleich oder doppelt so groß wie die kleinste Dimension ist. Für Beispiel: 4x4x8, 4×8×8 oder 12x12x24.

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

Die Kanten in diesem Topologiediagramm werden zur Verdeutlichung als ungerichtete Kanten dargestellt. In ist jede Kante 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 entsteht ein vollständig symmetrischer 4 × 2-verdrillter Torus. Topologie:

Zwischen diesem und dem vorherigen Diagramm hat sich nur das Y umlaufenden Kanten. 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.

Weitere Informationen finden Sie unter AcceleratorConfig verwenden. zum Angeben einer Twisted-Tor-Konfiguration beim Erstellen einer Cloud TPU.

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

Verdrehte Topologie	Theoretische Erhöhung der Teilung Bandbreite im Vergleich zu einem nicht verdrehten 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. Um z. B. eine gedrehte Form von 4x4x8 anzugeben, würden Sie das Flag topology so angeben: 4x4x8_twisted.

TPU v4-Topologievarianten

Einige Topologien mit derselben Anzahl von Chips können auf unterschiedliche Weise angeordnet werden. Ein TPU-Pod-Slice mit 512 Chips (1.024 TensorCores) kann beispielsweise Topologien 4x4x32, 4x8x16 oder 8x8x8 konfiguriert. Ein TPU-Pod-Slice mit 2.048 Chips (4.096 TensorCores) bietet noch mehr Topologieoptionen: 4x4x128, 4 x 8 x 64, 4 x 16 x 32 und 8 x 16 x 16.

Die Standardtopologie, die einer bestimmten Chipanzahl zugeordnet ist, ist die Topologie, ähnlich einem Würfel. Diese Form ist wahrscheinlich die beste Wahl für das datenparrallele ML-Training. Sonstiges Topologien können für Arbeitslasten mit mehreren Arten von Parallelität (für Beispiel: 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 für die X‑Achse und eine 256‑Wege-Datenparallelität für die Y‑ und Z‑Achse festlegen, 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). Beispiel: Für eine TPU v4 Pod-Slice mit einer Topologie von 8 x 16 x 16. Die TPU-Topologieabmessungen sind 8, 16 und 16. Es ist 8- oder 16-Wege-Modellparallelität (die einer der Methoden physische TPU-Topologiedimensionen verwendet werden. Eine 4-Wege-Modellparallelität wäre nicht optimal. da sie nicht an eine der TPU-Topologiedimensionen ausgerichtet ist, aber mit einer 4x16x32-Topologie mit derselben Anzahl von Chips wäre es optimal.

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, die kleiner als 64 Chips sind: ein Würfel im Format 4 x 4 x 4. Sie können diese kleinen v4-Topologien entweder TensorCore-basierter Name (z. B. v4-32) oder deren Topologie (z. B. 2x2x4):

Name (basierend auf der TensorCore-Anzahl)	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. In der folgenden Tabelle sind einige Beispiele aufgeführt. 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 TensorCore-Anzahl)	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: muss die Flags --type und --topology verwenden	1.024	4x16x16
v4-4096	2.048	8x16x16
…	…	…