TPU v4

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

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

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)
Bandbreite pro Pod insgesamt reduzieren	1,1 PB/s
Biabschnitte Bandbreite pro Pod	24 TB/s

Das folgende Diagramm zeigt einen TPU v4-Chip.

3D-Mesh und 3D-Torus

v4-TPUs haben eine direkte Verbindung zu den nächsten benachbarten Chips in 3 So entsteht ein 3D-Mesh aus Netzwerkverbindungen. Die Verbindungen kann als 3D-Torus auf Slices konfiguriert werden, in denen die Topologie AxBxC entweder 2A=B=C oder 2A=2B=C, 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

NUMA (Non Uniform Memory Access) ist eine Computerspeicherarchitektur für Maschinen die mehrere CPUs haben. Jede CPU hat direkten Zugriff auf einen Block mit Hochgeschwindigkeitssystemen, zu speichern. Eine CPU und ihr Arbeitsspeicher werden als NUMA-Knoten bezeichnet. NUMA-Knoten sind mit anderen NUMA-Knoten, 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 die von einer CPU benötigten Daten seinen NUMA-Knoten, wodurch der Arbeitsspeicherdurchsatz erhöht wird. Weitere Informationen zu NUMA Siehe Non Uniform Memory Access. auf Wikipedia.

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

So aktivieren Sie die NUMA-Knotenbindung:

1. Installieren Sie das numactl-Befehlszeilentool.

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

2. Binden Sie den Skriptcode an NUMA Node 0.

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

Aktivieren Sie die NUMA-Knotenbindung, wenn:

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

Weitere Unterschiede im Speichersystem:

• v4-TPU-Chips haben auf dem gesamten Chip einen einheitlichen Arbeitsspeicher von 32 GiB HBM. Dies ermöglicht eine bessere Koordination zwischen den beiden On-Chip-TensorCores.
• Verbesserte HBM-Leistung dank neuer Speicherstandards und -geschwindigkeiten.
• Verbessertes DMA-Leistungsprofil mit integrierter Unterstützung für hohe Leistung mit 512 Mrd. Detaillierungsgraden.

TensorCores

• Die doppelte Anzahl an MXUs und eine höhere Taktrate mit maximal 275 TFLOPS-Werten.
• Doppelte Bandbreite für Transposition und Permutation.
• Load-Store-Arbeitsspeicherzugriffsmodell für Common Memory (Cmem).
• Schnelleres Laden der MXU-Gewichtsbandbreite und 8-Bit-Modus für niedrigere und die Inferenzlatenz verbessern.

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 TPU v4 Pod-Slice erstellen, können Sie seinen Typ und seine Größe auf zwei Arten angeben: AcceleratorType und AccleratorConfig.

Hinweis: AcceleratorType und AccleratorConfig schließen sich gegenseitig aus. können Sie Verwenden Sie beim Erstellen einer TPU nur einen.

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 enthält und die Anzahl der TensorCores, die Sie verwenden möchten. Um beispielsweise eine v4-Pod-Slice mit 32 TensorCores würden Sie --accelerator-type=v4-32 verwenden.

Der folgende Befehl erstellt ein v4-TPU-Pod-Slice mit 512 TensorCores unter Verwendung von Das Flag --accelerator-type:

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

Zum Konfigurieren von v4-TPUs mit AcceleratorConfig verwenden Sie die --version und die --topology-Flags. 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 3-Tupel AxBxC an, wobei A<=B<=C und A, B, C sind entweder alle <= 4 oder alle ganzzahligen Vielfachen von 4 sind. Die Werte A, B und C sind zählt der Chip in jeder der 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. Dies wird als Twisted-tori-Topologie bezeichnet.

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

Größere Pod-Slices können aus einem oder mehreren 4x4x4-„Cubes“ erstellt werden von Chips.

Verdrehte Tori-Topologien

Bei einigen 3D-Torus-Slice-Formen in Version 4 kann eine sogenannte Twisted Torus-Topologie. Zwei V4-Cubes können beispielsweise als 4 x 4 x 8 angeordnet werden. Slice oder 4x4x8_twisted. 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.

Für Arbeitslasten, die Datenparallelität als einzige Parallelitätsstrategie verwenden, Topologien funktionieren möglicherweise etwas besser. Für LLMs ist die Leistung mit einem Twisted Topologie kann je nach Art der Parallelität (DP, MP usw.) variieren. Beste trainieren Sie Ihr LLM mit und ohne verdrillte Topologie, die die beste Leistung für Ihr Modell bietet. Einige Experimente zur FSDP MaxText-Modell wurden mit einer Twisted-Topologie ein bis zwei Verbesserungen bei der Multi-Faktor-Authentifizierung festgestellt.

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

• Verbessertes Load-Balancing
• Höhere geteilte 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 Ränder in diesem Topologiediagramm werden zur besseren Übersichtlichkeit als ungerichtete Kanten dargestellt. In ist jede Kante eine bidirektionale Verbindung zwischen TPUs. Wir verweisen auf die und Kanten zwischen einer Seite des Rasters und der anderen Seite als umlaufende Kanten, wie in der Abbildung 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 demselben X -Koordinate wurden sie so verschoben, dass sie mit Koordinaten eine Verbindung zur TPU herstellen. X+2 Mod 4.

Dieselbe Idee verallgemeinert unterschiedliche Dimensionsgrößen und unterschiedliche Dimensionen. Das resultierende Netzwerk ist symmetrisch, solange jede Dimension der kleinsten Dimension gleich oder doppelt so groß ist.

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

Die folgende Tabelle zeigt die unterstützten verdrillten Topologien und eine theoretische mit ihnen gegenüber Topologien, bei denen dies nicht der Fall ist.

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 Twisted-Topologie wird nur dann automatisch verwendet, wenn das Suffix _twisted wird mit dem Flag --topology angegeben. Um z. B. eine gedrehte Form von 4x4x8 würden Sie das Flag topology so angeben: 4x4x8_twisted.

TPU v4-Topologievarianten

Einige Topologien mit der gleichen Anzahl von Chips können in unterschiedlichen . 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 datenparalleles 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 funktionieren am besten, wenn die Topologie auf die verwendete Parallelität abgestimmt ist. Wenn Sie z. B. die 4-Wege-Modellparallelität auf die X-Dimension und die 256-Wege-Daten anwenden, Parallelität der Y- und Z-Dimensionen entspricht einer 4x16x16-Topologie.

Modelle mit mehreren parallelen Dimensionen funktionieren am besten mit ihren Parallelitätsdimensionen, die TPU-Topologiedimensionen zugeordnet sind. Das sind in der Regel Data and Model parallel 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, deren Topologien kleiner sind 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-Slices sind in Schritten von 64 Chips verfügbar, wobei Formen ein Vielfaches von 4 für alle drei Dimensionen. Die Abmessungen müssen in aufsteigender Reihenfolge. In der folgenden Tabelle sind einige Beispiele aufgeführt. Einige der Diese Topologien sind „benutzerdefiniert“ Topologien, die nur mit den Die Flags --type und --topology, da es mehr als eine Möglichkeit gibt, das Element anzuordnen die Chips.

Name (basierend auf der TensorCore-Anzahl)	Anzahl der Chips	Topologie
v4-128	64	4x4x4
v4-256	128	4x4x8
v4-512	256	4x8x8
Nicht zutreffend – es müssen die Flags --type und --topology verwendet werden	256	4x4x16
v4-1024	512	8x8x8
v4-1536	768	8x8x12
v4-2048	1.024	8x8x16
Nicht zutreffend – es müssen die Flags --type und --topology verwendet werden	1.024	4x16x16
v4-4096	2.048	8x16x16
…	…	…