Systemarchitektur
Tensor Processing Units (TPUs) sind von Google entwickelte ML-Beschleuniger. Cloud TPU macht TPUs als skalierbare Google Cloud-Ressource verfügbar. Sie können Arbeitslasten für maschinelles Lernen auf Cloud TPUs mit ML-Frameworks wie TensorFlow, Pytorch und JAX ausführen.
Jeder TensorCore auf einem TPU-Chip besteht aus einer oder mehreren Matrix-Multiply-Einheiten (MXUs), einer Vektoreinheit und einer Skalareinheit.
Eine MXU setzt sich aus 128 x 128-Multiplikations-/Akkumulatoren in einem systolischen Array zusammen. Die MXUs stellen den Großteil der Rechenleistung in einem TensorCore bereit. Jede MXU kann in jedem Zyklus 16.000 Multiplikations- bzw. Akkumulationsvorgänge ausführen. Alle Faktoren nehmen bfloat16-Eingaben vor, aber alle Akkumulationen werden im FP32-Zahlenformat ausgeführt.
Die Vektoreinheit wird für allgemeine Berechnungen wie Aktivierungen und Softmax verwendet. Die Skalareinheit wird für die Ablaufsteuerung, die Berechnung der Speicheradresse und andere Wartungsvorgänge genutzt.
TPU-Versionen
Das genaue Layout einer TPU hängt von der verwendeten TPU-Version ab. Architekturdetails und Leistungsmerkmale von TPU v2 und v3 sind unter Domainspezifischer Supercomputer zum Trainieren neuronaler Netzwerke verfügbar.
TPU v4
Jeder v4 TPU-Chip enthält zwei TensorCores. Jeder TensorCore hat vier MXUs, eine Vektoreinheit und eine Skalareinheit. In der folgenden Tabelle sind die wichtigsten Spezifikationen und ihre Werte für einen v4 TPU-Pod aufgeführt.
| Wichtige Spezifikationen | v4 Pod-Werte |
|---|---|
| Spitzenrechenzeit 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 | 4096 Chips |
| Interconnect-Topologie | 3D Torus |
| Spitzenrechenleistung pro Pod | 1.1 Exaflops (bf16 oder int8) |
| Geringere Bandbreite pro Pod | 1,1 PB/s |
| Bisection-Bandbreite pro Pod | 24 TB/s |
Das folgende Diagramm zeigt einen TPU v4-Chip.
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 | Pod-Werte von v3 |
|---|---|
| Spitzenrechenzeit 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 |
| Spitzenrechenleistung pro Pod | 126 PetaFLOPS (bf16) |
| Geringere Bandbreite pro Pod | 340 TB/s |
| Bisection-Bandbreite pro Pod | 6,4 TB/s |
Das folgende Diagramm zeigt einen TPU v3-Chip.
TPU v2
Die kleinste TPU v2-Konfiguration enthält vier TPU-Chips und 16 GiB HBM. Jeder TPU-Chip enthält zwei TensorCores. Jeder TensorCore hat eine MXU, eine Vektoreinheit und eine Skalareinheit. Das folgende Diagramm zeigt einen TPU v2-Chip.
Cloud TPU bietet die folgenden TPU-Konfigurationen:
- Single-Board-TPU
- TPU-Pod – eine Gruppe von TPU-Geräten, die über Hochgeschwindigkeitsverbindungen verbunden sind
- TPU-Slice – eine Unterteilung eines TPU-Pods
Leistungsvorteile von TPU v4 gegenüber v3
Speichersystem:
- v4-TPU-Chips haben einen einheitlichen 32-GiB-HBM-Speicher auf dem gesamten Chip, der eine bessere Koordination zwischen den beiden On-Chip-TensorCores ermöglicht.
- Verbesserte HBM-Leistung mit den neuesten Speicherstandards und -geschwindigkeiten.
- Verbessertes DMA-Leistungsprofil mit integrierter Unterstützung für hohe Leistung bei 512B Granularitäten.
TensorCores:
- Doppelt so viele MXUs und eine höhere Taktrate, die maximal 275 TFLOPS liefern.
- 2-fache Transpositions- und Permutationsbandbreite.
- Speicherzugriffsmodell für Speicher für gemeinsamen Speicher (Cmem).
- Schnellere MXU-Gewichtungsbandbreite und Unterstützung des 8-Bit-Modus, um geringere Batchgrößen zu ermöglichen und die Latenz der Inferenz zu verbessern.
Inter-Chip-Interconnect:
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 größere Anzahl von FLOPS pro Kern und größere Speicherkapazität bei TPU v3-Konfigurationen können die Leistung von Modellen auf folgende Weise verbessern:
TPU v3-Konfigurationen bieten für rechengebundene Modelle erhebliche Leistungsvorteile pro Kern. Für speichergebundene Modelle kann im Vergleich zu TPU v2-Konfigurationen möglicherweise nicht dieselbe Leistungsverbesserung erzielt werden, wenn sie auch bei TPU v3-Konfigurationen speichergebunden sind.
In Fällen, in denen Daten in TPU v2-Konfigurationen nicht in den Arbeitsspeicher passen, kann TPU v3 eine bessere Leistung und eine geringere Neuberechnung von Zwischenwerten bieten (Rematerialisierung).
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.
TPU v4-Konfigurationen
Ein TPU v4 Pod besteht aus 4.096 Chips, die mit neu konfigurierbaren Hochgeschwindigkeitsverbindungen verbunden sind. Mit dem flexiblen Netzwerk von TPU v4 können Sie die Chips auf verschiedene Arten in einem Segment derselben Größe verbinden. Sie können v4-TPUs mit AcceleratorConfig konfigurieren.
Mit AcceleratorConfig können Sie die Segmentgröße in Form von Chips über das Feld topology angeben.
Die TPU v4-Topologie wird mit einem 3-Tupel angegeben, das die Anzahl, die Form und die Verbindungen zwischen TPU-Chips beschreibt. Die folgenden Abbildungen zeigen einige gängige TPU v4-Topologien.
Größere Segmente können aus einem oder mehreren 4x4x4-Würfeln aus Chips gebaut werden.
TPU-Segmente einer bestimmten Anzahl von Chips können auf verschiedene Arten konfiguriert werden. Ein TPU-Slice mit dem AcceleratorType v4-1024 kann beispielsweise so konfiguriert werden: 4x4x32, 4x8x16 oder 8x8x8.
TPU v2- und v3-Konfigurationen
TPUs sind in folgenden Konfigurationen verfügbar:
- Einzelnes TPU-Board
- TPU-Pod
- TPU-Pod-Slice
Single-Board-TPU
Eine Single-Board-TPU-Konfiguration ist ein eigenständiges Board mit vier TPU-Chips (acht TensorCores) ohne Netzwerkverbindungen mit anderen TPU-Boards. Single-Board-TPUs sind nicht Teil einer TPU-Pod-Konfiguration und nehmen keinen Teil eines TPU-Pods für sich in Anspruch.
TPU-Pods und -Slices
In einem TPU Pod- oder TPU Pod-Slice werden TPU-Chips über eine Hochgeschwindigkeitsverbindung miteinander verbunden. Jeder TPU-Chip kommuniziert direkt mit den anderen Chips auf dem TPU-Gerät. Die TPU-Software verwaltet automatisch die Verteilung von Daten an jeden TPU-Kern in einem Pod oder Slice. Pod-Slices sind mit 32, 128, 512, 1.024 oder 2.048 Kernen verfügbar.
Cloud TPU-VM-Architekturen
TPUs sind darauf ausgelegt, Matrixvorgänge schnell auszuführen. Jedes TPU-Board ist mit einer CPU-basierten Hostmaschine verbunden, um Vorgänge auszuführen, die nicht auf der TPU ausgeführt werden können. Die Hostcomputer sind für das Laden von Daten aus Cloud Storage, das Vorverarbeiten von Daten und das Senden von Daten an die TPU zuständig.
In einem TPU-Pod gibt es für jedes TPU-Board einen TPU-Host.
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
Die TPU-Knotenarchitektur besteht aus einer Nutzer-VM, die über gRPC mit dem TPU-Host kommuniziert. Wenn Sie diese Architektur verwenden, können Sie nicht direkt auf den TPU-Host zugreifen, was das Debuggen von Trainings und TPU-Fehlern erschwert.
TPU-VMs
Mit der TPU-VM-Architektur können Sie eine SSH-Verbindung zur VM herstellen, die physisch mit dem TPU-Gerät verbunden ist. Sie haben Root-Zugriff auf die VM, sodass Sie beliebigen Code ausführen können. Sie können auf Compiler- und Laufzeit-Debugging-Logs und Fehlermeldungen zugreifen.
Unterstützte Frameworks
Frameworks wie JAX, PyTorch und TensorFlow greifen über eine gemeinsam genutzte Bibliothek namens libtpu auf TPUs zu. libtpu enthält den XLA-Compiler, die TPU-Laufzeitsoftware und den TPU-Treiber.
TensorFlow
TPU-VMs
Mit TPU-VMs kann Ihr Python-Code statt auf einer Nutzer-VM direkt auf dem TPU-Host ausgeführt werden.
Weitere Informationen zu TensorFlow und Cloud TPU-VMs finden Sie im Nutzerhandbuch für Cloud TPU-VMs.
TPU-Knoten
Die Cloud TPU-Knotensystemarchitektur wurde ursprünglich für TensorFlow entwickelt. Auf die TPU-Hosts kann der Nutzer nicht zugreifen und keine monitorlose Kopie des TensorFlow-Servers ausführen. Die Hosts führen weder Python noch Nutzercode aus, der nicht als TensorFlow-Diagramm dargestellt wird. Nutzercode wird in einer separaten Nutzer-VM ausgeführt, die über das gRPC-Netzwerk mit den TPU-Hosts kommuniziert.
PyTorch
TPU-VMs
Bei TPU-VMs wird Ihr PyTorch-Code direkt auf den TPU-Hosts ausgeführt.
Weitere Informationen zu PyTorch und Cloud TPU finden Sie im Nutzerhandbuch für PyTorch/XLA.
TPU-Knoten
PyTorch wird in der Cloud TPU-Knotenarchitektur mit einer Bibliothek namens XRT ausgeführt. XRT sendet XLA-Diagramme und Laufzeitanweisungen über gRPC zur Ausführung auf TensorFlow-Servern. Für jeden TPU-Host ist eine Nutzer-VM erforderlich.
JAX
TPU-VMs
Bei TPU-VMs wird Ihr JAX-Code direkt auf den TPU-Hosts ausgeführt.
Weitere Informationen zum Ausführen von JAX auf Cloud TPU finden Sie unter JAX-Kurzanleitung.
TPU-Knoten
JAX auf Cloud TPU-Knoten wird wie mit PyTorch ausgeführt, da für jede Host-VM eine separate Nutzer-VM erforderlich ist.
