- TPU ist ein von Google entwickelter maßgeschneiderter Chip für großskaliges KI-Training und Inferenz und verfolgt eine andere Designphilosophie als GPUs
- Der Fokus liegt auf Skalierbarkeit und Energieeffizienz; Hardware (z. B. System-on-Chip-Aufbau, großer On-Chip-Speicher) und Software (XLA-Compiler) werden gemeinsam entworfen
- Die Kernarchitektur basiert auf systolischen Arrays, Pipelining und einem Ahead-of-Time-Kompilierungsansatz und ist für die meisten Deep-Learning-Operationen (insbesondere Matrixmultiplikation) optimiert
- Dank OCI- und OCS-Technologien sind flexible Knotenkonfigurationen, leistungsstarke Parallelverarbeitung und verschiedene Topologieoptionen möglich
- TPU-Systeme lassen sich hierarchisch von einem einzelnen Chip bis zu riesigen Multipods skalieren, was das Training großer KI-Modelle und eine effiziente Ressourcennutzung ermöglicht
Überblick und Hintergrund
- TPU ist ein auf ASIC basierender Chip, den Google entwickelt hat, um auf die massive Skalierung von KI-Diensten zu reagieren
- Anfangs wurden GPU, FPGA und ASIC gegeneinander abgewogen, doch mit der Ausweitung von Deep-Learning-Anwendungen wie der Sprachsuche im Jahr 2013 stieg der Bedarf an spezialisierter Hardware
- Heute spielt TPU eine unverzichtbare Rolle in wichtigen KI-Diensten wie Gemini und Veo und wird auch breit in Empfehlungsmodellen wie DLRM eingesetzt
Aufbau eines einzelnen TPU-Chips
Grundkonfiguration
- Beim TPUv4 verfügt ein einzelner Chip über 2 TensorCores (bei inferenzspezialisierten TPUs 1)
- Jeder TensorCore ist mit CMEM- (128 MiB) und HBM-Speicherbausteinen (32 GiB) verbunden
Interner Aufbau des TensorCore
- Matrix Multiply Unit (MXU): basiert auf einem systolischen 128x128-Array und ist ausschließlich für Matrixmultiplikationen zuständig
- Vector Unit (VPU): führt allgemeine elementweise Operationen aus
- Vector Memory (VMEM; 32 MiB): kopiert Daten aus dem HBM und dient als vorbereitender Arbeitsspeicher
- Scalar Unit + Scalar Memory (SMEM; 10 MiB): zuständig für Kontrollfluss, skalare Berechnungen und Speicheradressverwaltung
Strukturelle Unterschiede zu GPUs
- TPUs haben deutlich größeren On-Chip-Speicher (CMEM, VMEM, SMEM) als GPUs
- Die HBM-Kapazität ist bei GPUs größer, und auch die Zahl der Rechenkerne ist dort wesentlich höher
- Beim TPUv5p werden 500 TFLOPs/sec pro Chip und 4.45 ExaFLOPs/sec für ein gesamtes Pod (8960 Chips) erreicht
TPU-Designphilosophie
1. Systolische Arrays und Pipelining
- In einem systolischen Array sind Verarbeitungselemente (PEs) angeordnet, die ihre Rechenergebnisse an benachbarte Elemente weitergeben
- Nach der Dateneingabe sind kontinuierliche Berechnungen ohne zusätzliche Steuerung möglich; Speicherzugriffe zum Lesen und Schreiben finden nur bei Ein- und Ausgabe statt
- Die Architektur ist vor allem für Matrixmultiplikation und Convolution optimiert
- Durch Pipeline-Verarbeitung werden Rechnen und Datenbewegung gleichzeitig ausgeführt, um den Durchsatz zu maximieren
Nachteil systolischer Arrays – Sparsity
- Da alle Rechenelemente ständig aktiv sind, eignen sie sich nicht gut für sparse Matrizen
- Wenn DL-Modelle künftig stärker auf unregelmäßige Sparsity setzen, kann das zu einer Einschränkung werden
2. Ahead-of-Time-Kompilierung (AoT) und minimale Cache-Nutzung
- Durch das TPU-XLA-Codesign wird die Abhängigkeit von Caches, die unregelmäßige Speicherzugriffe benötigen, minimiert und so der Energieverbrauch gesenkt
- Der XLA-Compiler analysiert den Rechengraphen und berechnet Speicherzugriffsmuster im Voraus; statt auf Caches setzt das System primär auf Scratchpad-Speicher
- Das
@jit von JAX ist eine Zwischenform aus JIT und AoT: Beim ersten Lauf wird ein statischer Graph erzeugt und anschließend per XLA AoT-kompiliert
- Wenn sich die Eingabeform ändert, ist eine Neukompilierung nötig; bei dynamischem Padding oder Schleifen ist das ineffizient
- Dadurch wird eine hohe Energieeffizienz erreicht, allerdings auf Kosten der Flexibilität
Energieeffizienz von TPUv4
- Moderne Chips verwenden HBM3-Speicher, um den Energieverbrauch zu senken
- Da Speicheroperationen im Vergleich zu Rechenoperationen um ein Vielfaches mehr Energie verbrauchen, kann die Minimierung von Speicherzugriffen die Effizienz stark verbessern
TPU-Architektur mit mehreren Chips
Tray-Ebene (Tray/Board; 4 Chips)
- Ein Tray besteht aus 4 TPU-Chips (8 TensorCores) und einem CPU-Host
- Die Verbindung Host↔Chip erfolgt über PCIe, Chip↔Chip über Inter-Core Interconnect (ICI), das eine höhere Bandbreite bietet
Rack-Ebene (Rack; 4x4x4 = 64 Chips)
- Ein Rack besteht aus 64 TPU-Chips und ist über ICI und OCS (Optical Circuit Switching) als 3D-Torus (4x4x4) verbunden
- Google unterscheidet zwischen Rack, Pod und Slice
- Rack: physische Einheit aus 64 Chips (= Würfel)
- Pod: größte Einheit, die über ICI und OCS verbunden werden kann (z. B. TPUv4 = 4096 Chips = 64 Racks)
- Slice: abstrakte Einheit für beliebige Konfigurationen zwischen 4 Chips und Superpod
Vorteile von OCS
- Wraparound: Jede Achse wird als Ring (1D-Torus) umgesetzt, wodurch sich die maximale Anzahl an Hops zwischen Knoten reduziert
- Flexible nicht zusammenhängende Slices: Durch die OCS-Switching-Struktur können auch physisch getrennte Knoten zu einem Slice zusammengestellt werden, was Ressourcennutzung und Wartung erleichtert
- Verdrehte Topologien: Bei gleicher Chipzahl (fixes x, y, z) lässt sich die Verbindungsstruktur ändern, um die Geschwindigkeit für bestimmte Rechenmuster zu optimieren (z. B. twisted torus)
Beispiele für den Einsatz von Topologien
- Würfelform: geeignet für Daten- und Tensor-Parallelität (maximale Bandbreite)
- Lineare Form (Zigarrenform): vorteilhaft für Pipeline-Parallelität
- Twisted Torus: beschleunigt all-to-all-Kommunikation (z. B. bei Tensor Parallelism)
Superpod (Full Pod/Superpod; TPUv4: 4096 Chips / 64 Racks)
- Mehrere Racks werden über ICI und OCS zu einem extrem großen System verbunden
- Je nach Slice-Topologie unterscheiden sich Kommunikationsbandbreite und Parallelverarbeitungsleistung
- Dank OCS werden flexible Konfigurationen wie nicht zusammenhängende Slices oder verdrehte Topologien möglich
Multipod-Ebene (Multi-pod/Multislice; TPUv4: mehr als 4096)
- Mehrere Pods werden über das Datacenter Network (DCN) verbunden, was großskalige Trainingsinfrastruktur ermöglicht, allerdings mit geringerer Bandbreite als ICI
- Das Training von PaLM nutzte 2 Pods (6144 TPUv4) und wurde mit Ressourcenmanagement über insgesamt 6 Pods durchgeführt
- Bei der Entwicklung großer Modelle optimiert der XLA-Compiler Kommunikationsmuster anhand von Forschereinstellungen wie Parallelitätsdimensionen
- XLA fügt Kommunikationsoperationen zwischen einzelnen Slices und Pods ein und ermöglicht so großskaliges verteiltes Training mit minimalen Codeänderungen
Beispiele für reale Hardware und Diagramme
- TPU-Rack: Ein 4x4x4-3D-Torus bildet eine Einheit, jede Reihe enthält 2 Trays (8 Chips)
- TPUv4-Tray: In der Praxis gibt es 4 PCIe-Ports (jeweils einer pro TPU)
- TPUv4-Chip: zentral das ASIC, darum herum 4 HBM-Stacks (bei 2 TensorCores)
- TPUv4i-Chip-Floorplan (Inferenzvariante): 1 TensorCore, große CMEM-Fläche
Abschluss
- Dank an Google TPU Research Cloud (TRC) für die Forschungsunterstützung
Referenzen
Noch keine Kommentare.