Systemarchitektur

Tensor Processing Units (TPUs) sind anwendungsspezifische integrierte Schaltungen (ASICs), die von Google entwickelt wurden, um Arbeitslasten beim maschinellen Lernen zu beschleunigen. Cloud TPU ist ein Google Cloud-Dienst, der TPUs als skalierbare Ressource zur Verfügung stellt.

TPUs sind darauf ausgelegt, Matrixvorgänge schnell durchzuführen. Sie eignen sich daher ideal für Arbeitslasten beim maschinellen Lernen. Sie können ML-Arbeitslasten auf TPUs mit Frameworks wie TensorFlow, Pytorch und JAX ausführen.

Cloud TPU-Nutzungsbedingungen

In den folgenden Abschnitten werden die in diesem Dokument verwendeten Begriffe und verwandte Konzepte erläutert.

TPU-Chip

Ein TPU-Chip enthält einen oder mehrere TensorCores. Jeder TensorCore besteht aus einer oder mehreren Matrix-Multiply-Einheiten (MXUs), einer Vektoreinheit und einer skalaren Einheit.

Eine MXU besteht aus 128 x 128 Multiplikatoren/Akkumulatoren in einem systolic Array. MXUs bieten den Großteil der Rechenleistung in einem TensorCore. Jede MXU kann 16.000 Multiplikations- bzw. Akkumulationsvorgänge pro Zyklus ausführen. Alle Multiplikationen akzeptieren bfloat16-Eingaben, aber alle Akkumulierungen werden im FP32-Zahlenformat ausgeführt.

Die Vektoreinheit wird für allgemeine Berechnungen wie Aktivierungen und Softmax-Funktionen verwendet. Die Skalareinheit wird für die Ablaufsteuerung, die Berechnung von Speicheradressen und andere Wartungsvorgänge verwendet.

TPU-Gerät

Ein TPU-Gerät enthält vier TPU-Chips, die über Hochgeschwindigkeitsverbindungen miteinander verbunden sind. Die Anzahl der TensorCores in einem TPU-Chip und die Anzahl der MXUs in einem TensorCore hängt von der Version des verwendeten TPU-Geräts ab.

Hostcomputer

Jedes TPU-Gerät ist mit einer CPU-basierten Hostmaschine (TPU-VM) verbunden. Der Hostcomputer ist dafür verantwortlich, Daten aus Cloud Storage zu laden, Daten vorzuverarbeiten, Daten an die TPU zu senden und Vorgänge auszuführen, die nicht auf einer TPU ausgeführt werden können.

Einzelne Geräte und Pods

Mit Cloud TPU können Sie mit einem einzelnen TPU-Gerät oder mehreren miteinander verbundenen TPU-Geräten, sogenannten Pods, arbeiten. Sie können mit einem ganzen TPU-Pod arbeiten (die Anzahl der TPU-Geräte in einem Pod variiert je nach TPU-Version) oder einer Teilmenge eines TPU-Pods, die als Segment bezeichnet wird.

TPU-Einzelgeräte

Eine TPU-Konfiguration für ein einzelnes Gerät ist ein eigenständiges Board ohne Netzwerkverbindungen zu anderen TPU-Boards. Single-Board-TPUs sind nicht Teil einer TPU-Pod-Konfiguration und können nicht mit anderen TPU-Geräten kommunizieren.

TPU-Segmente

In einem TPU-Slice sind die TPU-Chips über eine Hochgeschwindigkeits-Interconnect-Verbindung verbunden. Jeder TPU-Chip kommuniziert direkt mit den anderen Chips im TPU-Segment. Die TPU-Software übernimmt automatisch die Verteilung von Daten an jeden TensorCore in einem Segment. TPU v4-Segmente sind mit 16, 32, 64, 128, 256, 512, 1024, 2048 oder 4096 TensorCores verfügbar. TPU v3-Segmente sind mit 32, 128, 512, 1.024 oder 2.048 TensorCores verfügbar. TPU v2-Segmente sind mit 32, 128, 256 oder 512 TensorCores verfügbar.

TPU-Versionen

Das genaue Layout einer TPU hängt von der verwendeten TPU-Version ab. Architekturdetails und Leistungsmerkmale von TPU v2 und v3 finden Sie unter A Domain Specific Supercomputer for Training Deep Neural Networks.

TPU v4

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

Wichtige Spezifikationen	V4 Pod-Werte
Spitzenverarbeitung pro Chip	275 TeraFLOPS (bf16 oder int8)
HBM2-Kapazität und Bandbreite	32 GiB, 1.200 Gbit/s
Gemessene minimale/mittlere/maximale Leistung	90/170/192 W
TPU-Pod-Größe	4.096 Chips
Interconnect-Topologie	3D-Mesh-Netzwerk
Spitzen-Computing pro Pod	1.1 Exaflops (bf16 oder int8)
Gesamte Bandbreite pro Pod reduzieren	1,1 PB/s
Bandbreite pro Pod aufteilen	24 TB/s

Das folgende Diagramm veranschaulicht einen TPU v4-Chip.

3D-Mesh-Netzwerk und 3D-Torus

v4-TPUs haben in drei Dimensionen eine direkte Verbindung zu den nächsten benachbarten Chips, was zu einem 3D-Mesh von Netzwerkverbindungen führt. Wenn das Segment größer oder gleich einem einzelnen Würfel ist, können die Verbindungen als 3D-Torus konfiguriert werden. Wenn alle Segmentgrößen ein vollständiger 3D-Torus sind (Wrap-around-Links verbunden), ist die Leistung besser als mit einer 3D-Mesh-Konfiguration.

v4-Topologie

Einige 3D-Torus-Slice-Formen in Version 4 haben die Möglichkeit, eine sogenannte Twisted Torus-Topologie zu verwenden, um die Querschnittsbandbreite des Segments zu erhöhen. Eine Twisted Topologie ist eine alternative Topologie für v4-TPUs.

Hinweis: Die Twisted-Topologie ist nur für 3D-Torus-TPU-Konfigurationen und somit nicht für v2- oder v3-TPU-Konfigurationen verfügbar, die 2D-Mesh-Topologien verwenden.

Einige v4-Topologien bieten die Möglichkeit einer verdrillten Konfiguration. Beispielsweise können zwei v4-Cubes als 4x4x8-Slice oder 4x4x8_twisted angeordnet werden. Twisted Tori bietet eine erheblich höhere Bandbreite in Sektionen. Von dieser erhöhten Bandbreite können Arbeitslasten mit globalen Kommunikationsmustern profitieren. Es wurde festgestellt, dass verdrehte Topologien die Leistung der meisten Modelle verbessern, wobei große Einbettungsmodelle am stärksten davon profitieren.

Die Twisted Tori-Variante wird für das LLM-Training (Large Language Model) empfohlen, das Untergruppen umfasst. Twisted Tori wird wahrscheinlich am meisten von Einbettungsarbeitslasten profitieren. Bei Arbeitslasten mit reiner Datenparallelität kann mit Twisted Tori eine etwas bessere Leistung erzielt werden. Bei LLMs mit mehreren Formen der Parallelität (DP, MP usw.) hat ein Twist-Tori wahrscheinlich eine geringere Leistung als ein regulärer Torus.

Der Hauptvorteil des „Verdrehens“ eines Torus besteht darin, dass er eine asymmetrische Torus-Topologie (z. B. 4 × 4 × 8) in eine eng verwandte Topologie umwandelt, die symmetrisch ist. Die symmetrische Topologie hat viele Vorteile: besseres Load-Balancing, höhere Bandbreite in einzelnen Abschnitten und kürzere Paketrouten. Diese Vorteile führen letztendlich zu einer verbesserten Leistung für viele globale Kommunikationsmuster.

Die TPU-Software unterstützt Twisted Tori auf Segmenten, bei denen die Größe jeder Dimension entweder gleich oder doppelt so groß wie die kleinste Dimension ist, z. B. 4 × 4 × 8, 4 × 8 × 8 oder 12 × 12 × 24.

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

Die Kanten in diesem Topologiediagramm werden der Übersichtlichkeit halber als ungerichtete Kanten angezeigt. In der Praxis ist jede Edge eine bidirektionale Verbindung zwischen TPUs. Pakete können über diese Kanten in beide Richtungen übertragen werden. Die Kanten zwischen einer Seite dieses Rasters und der gegenüberliegenden Seite werden, wie im Diagramm dargestellt, als umlaufende Kanten bezeichnet.

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

Zwischen diesem und dem vorherigen Diagramm haben sich nur die Y-Umlaufkanten geändert. Anstatt eine Verbindung zu einer anderen TPU mit derselben X-Koordinate herzustellen, wurden sie verschoben, um eine Verbindung zur TPU mit Koordinate X+2 Mod 4 herzustellen.

Dasselbe gilt für unterschiedliche Dimensionsgrößen und eine unterschiedliche Anzahl von Dimensionen. Das resultierende Netzwerk ist symmetrisch, solange jede Dimension gleich oder doppelt so groß wie die kleinste Dimension ist.

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

Die folgende Tabelle zeigt die unterstützten verdrillten Topologien und die theoretische Steigerung der Bandbreite in geteilten Abschnitten im Vergleich zu unverdrillten Topologien.

Twisted-Topologie	Theoretische Steigerung der Bandbreite im Vergleich zu einem nicht verdrehten Toru
4×4×8_twisted	~70%
8x8x16_twisted
12×12×24_twisted
4×8×8_twisted	~40%
8×16×16_twisted

TPU v3

Jeder v3 TPU-Chip enthält zwei TensorCores. Jeder TensorCore hat zwei MXUs, eine Vektoreinheit und eine Skalareinheit. In der folgenden Tabelle sind die wichtigsten Spezifikationen und ihre Werte für einen v3 TPU-Pod aufgeführt.

Wichtige Spezifikationen	Werte von v3-Pods
Spitzenverarbeitung pro Chip	123 TeraFLOPS (bf16)
HBM2-Kapazität und Bandbreite	32 GiB, 900 Gbit/s
Gemessene minimale/mittlere/maximale Leistung	123/220/262 W
TPU-Pod-Größe	1.024 Chips
Interconnect-Topologie	2D-Torus
Spitzen-Computing pro Pod	126 PetaFLOPS (bf16)
Gesamte Bandbreite pro Pod reduzieren	340 TB/s
Bandbreite pro Pod aufteilen	6,4 TB/s

Das folgende Diagramm zeigt einen TPU v3-Chip.

Leistungsvorteile von TPU v4 gegenüber v3

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 NUMA-Knoten verbunden, die direkt nebeneinander stehen. Eine CPU von einem NUMA-Knoten 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 einer Maschine mit mehreren CPUs ausgeführt wird, kann von einer CPU benötigte Daten in ihrem NUMA-Knoten platzieren, was den Arbeitsspeicherdurchsatz erhöht. Weitere Informationen zu NUMA finden Sie auf Wikipedia unter Non Uniform Memory Access.

Die TPU v4-Laufzeit ist nativ nicht NUMA-sensitive. Sie können jedoch 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. Verwenden Sie numactl --cpunodebind=0, wenn Sie das Trainingsskript starten. Dadurch wird der Skriptcode an NUMA Node 0 gebunden.

    $ $ 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ängig ist (z. B. Imageklassifizierung, Empfehlungsarbeitslasten).
• Wenn Sie eine TPU-Laufzeitversion ohne das Suffix -pod verwenden (z. B. tpu-vm-tf-2.10.0-v4).

Weitere Unterschiede im Speichersystem:

• v4-TPU-Chips haben einen einheitlichen Speicherbereich von 32 GiB HBM auf dem gesamten Chip und ermöglichen eine bessere Koordination zwischen den beiden On-Chip-TensorCores.
• Verbesserte HBM-Leistung dank aktueller Speicherstandards und Geschwindigkeiten
• Verbessertes DMA-Leistungsprofil mit integrierter Unterstützung von bis zu 512 Milliarden Detailtiefen.

TensorCores:

• Die doppelte Anzahl von MXUs und eine höhere Taktrate liefern maximal 275 TFLOPS.
• 2-fache Transposition und Permutationsbandbreite.
• Load-Store-Arbeitsspeicherzugriffsmodell für Common Memory (Cmem).
• Schnellere MXU-Load-Balancing-Bandbreite und Unterstützung des 8-Bit-Modus, um niedrigere Batchgrößen zu ermöglichen und die Inferenzlatenz zu verbessern.

Inter-Chip-Interconnect-Verbindung:

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

Sonstiger:

• x16-PCIE-gen3-Schnittstelle zum Host (direkte Verbindung).
• Verbessertes Sicherheitsmodell.
• Verbesserte Energieeffizienz

Leistungsvorteile von TPU v3 im Vergleich zu v2

Die höheren FLOPS pro TensorCore und die höhere Arbeitsspeicherkapazität in TPU v3-Konfigurationen können die Leistung Ihrer Modelle auf folgende Weise verbessern:

• TPU v3-Konfigurationen bieten erhebliche Leistungsvorteile pro TensorCore für rechengebundene Modelle. Speichergebundene Modelle in TPU v2-Konfigurationen erreichen diese Leistungsverbesserung möglicherweise nicht, wenn sie auch in TPU v3-Konfigurationen speichergebunden sind.

• Wenn die Daten bei TPU v2-Konfigurationen nicht in den Arbeitsspeicher passen, bietet TPU v3 eine verbesserte Leistung und reduzierte Neuberechnung von Zwischenwerten (Neuberechnung von Zwischenwerten).

• Mit TPU v3-Konfigurationen können neue Modelle mit Stapelgrößen ausgeführt werden, für die TPU v2-Konfigurationen nicht ausgereicht haben. Zum Beispiel lässt TPU v3 tiefere ResNets und größere Bilder mit RetinaNet zu.

Modelle, die auf TPU v2-Hardware nahezu eingabegebunden arbeiten ("Infeed"), da Trainingsschritte auf Eingaben warten müssen, arbeiten unter Umständen auch auf Cloud TPU v3-Hardware eingabegebunden. Mit dem Leitfaden zur Leistung von Pipelines können Sie Einspeisungsprobleme beheben.

Cloud TPU-VM-Architekturen

Wie Sie mit dem TPU-Host (und dem TPU-Board) interagieren, hängt von der verwendeten TPU-VM-Architektur ab: TPU-Knoten oder TPU-VMs.

TPU-Knoten

Wichtig: TPU v4 wird mit der TPU-Knotenarchitektur nicht unterstützt.

Die TPU-Knotenarchitektur besteht aus einer Nutzer-VM, die über gRPC mit dem TPU-Host kommuniziert. Bei Verwendung dieser Architektur können Sie nicht direkt auf den TPU-Host zugreifen, was die Fehlerbehebung von Trainings- und TPU-Fehlern erschwert.

TPU-VMs

Mit der TPU-VM-Architektur können Sie eine SSH-Verbindung zu der VM herstellen, die physisch mit dem TPU-Gerät verbunden ist. Sie haben Root-Zugriff auf die VM und können daher beliebigen Code ausführen. Sie können auf Compiler- und Laufzeit-Fehlerbehebungsprotokolle und Fehlermeldungen zugreifen.

Nächste Schritte

• Kurzanleitung
• Anleitungen