Künstliche Intelligenz: Architektur und Performance von Googles KI-Chip TPU
Seit 2015 betreibt Google in seinen Rechenzentren selbst entwickelte Spezial-Chips zur Beschleunigung von KI-Anwendungen mit neuronalen Netzen, die Tensor Processing Units (TPUs). Jetzt erfährt man endlich etwas über deren Architektur und Performance.
(Bild: Google)
Kurz nachdem die Google-Tochter DeepMind im März 2016 mit ihrer künstlichen Intelligenz AlphaGo in einem aufsehenerregenden Turnier den Top-Profi Lee Sedol im Strategiespiel Go geschlagen hat, gab Google auf der I/O 2016 bekannt, dass dabei selbst entwickelte Spezialchips beteiligt waren: die Tensor Processing Units (TPUs). Außer einem Foto, das im Wesentlichen einen großen Kühlkörper zeigte, rückte Google damals aber praktisch keine Informationen über diese Wunderchips heraus. Jetzt gewährt Google einen Blick unter den Kühlkörper und veröffentlicht eine Studie über die Performance der TPUs im Vergleich zu CPUs und GPUs.
Die TPU ist auf die Anwendung neuronaler Netze optimiert – wohlgemerkt nicht auf die Trainingsphase. Ein künstliches neuronales Netz besteht aus (heutzutage vielen) Schichten von Neuronen. Jedes Neuron ermittelt eine gewichtete Summe der Ausgabewerte der Neuronen in der vorigen Schicht. Daraus ergibt sich dann mittels einer nichtlinearen Aktivierungsfunktion der Ausgabewert des Neurons, der als Eingabe in die nächste Schicht einfließt. Multiplizieren und Aufaddieren sind also die Grundoperationen bei der Anwendung neuronaler Netze, und sie fallen in rauen Mengen an.
Training auf GPUs, Anwendung auf TPUs
Beim Trainieren eines neuronalen Netzes werden die Gewichte der Verbindungen zwischen den Neuronen ermittelt, und dafür braucht man Gleitkommaoperationen. Dies ist daher nach wie vor eine Domäne der GPUs, die auf massiv paralleles Gleitkomma-Rechnen optimiert sind. Bei der anschließenden Anwendung des trainierten Netzes genügt es aber, die Gewichte zu quantisieren und mit geringerer Genauigkeit zu rechnen, Ganzzahlen genügen.
Und hier liegt das Geheimnis von Rechenleistung und Energieeffizenz der TPU: Die Google-TPU enthält eine Matrix aus 256 × 256 Multiplizier/Aufaddier-Einheiten (Multiplier/Accumulator, MAC), die mit 8-bittigen Ganzzahlen arbeiten (im Blockdiagramm gelb). Bei einem Takt von 700 MHz erreichen diese eine theoretische Spitzenleistung von 92 Tera-Operationen pro Sekunde (65.536 Multiplikationen + 65.536 Additionen, das Ganze mal 700 MHz Taktfrequenz ergibt 91,75 TOPS/s). Bei Bedarf können die Recheneinheiten auch mit 8-bittigen Gewichten und 16-bittigen Aktivierungswerten arbeiten oder umgekehrt (bei halber Geschwindigkeit) oder beides 16-bittig, dann mit einem Viertel der Geschwindigkeit.
(Bild: Google )
Datenpumpen
24 % der Chipfläche nehmen diese Recheneinheiten ein, und der ganze Rest dient dazu, sie möglichst pausenlos mit Daten zu versorgen. Ein 24 MByte großer Unified Buffer (29 % der Chipfläche) speichert die Zwischenwerte der Berechnungen und ist über ein 256 Byte (!) breites Interface an die Rechenmatrix angebunden. Zwischen diesem Unified Buffer und dem Host überträgt ein DMA-Controller Daten über 16 PCIe-3-Lanes.
Für die (im Verlauf der Anwendung des neuronalen Netzes konstanten) Gewichtsdaten sind extern 8 GByte DDR3-Speicher angebunden. Das genügt, um die Gewichte vieler verschiedener Anwendungen vorzuhalten. Ein Weight FIFO liefert die für die konkrete Anwendung benötigten Gewichte mit 30 GByte/s an die Rechenmatrix.
Programmierung
Das Host-System programmiert die TPU über einen CISC-Befehlssatz, der nur eine Handvoll Instruktionen umfasst. In deren Zentrum steht die Matrix-Multiplikation – nichts anderes ist all das oben beschriebene Multiplizieren und Aufaddieren. Die restlichen Instruktionen transferieren Daten zwischen TPU und Host, laden Gewichte, berechnen die Aktivierungsfunktionen oder dienen der Synchronisation.
Diese Maschinensprache ist maßgeschneidert für Googles TensorFlow-Framework. Darin programmiert man KI-Algorithmen auf Hochsprachen-Niveau und beschreibt sie als Datenflüsse von Vektoren und Matrizen. Mit passenden Backends kann man dann solche TensorFlow-Programme auf CPUs oder GPUs ausführen – oder eben auch auf TPUs.
Performance
Google hat die TPUs bereits seit 2015 in seinen Rechenzentren im Produktiv-Einsatz. In seiner Studie In-Datacenter Performance Analysis of a Tensor Processing Unit nimmt das TPU-Team die Performance der TPU unter die Lupe und vergleicht sie mit der (im Jahr 2015) aktueller CPUs und GPUs (Intel Haswell und Nvidia K80). Für den Vergleich haben die Forscher sechs produktive Anwendungen neuronaler Netze herangezogen – eine davon ist eines der neuronalen Netze aus AlphaGo.
Im Ergebnis war die TPU um den Faktor 15 bis 30 schneller als CPU und GPU. Noch besser sieht der Vergleich der Energieeffizienz aus: In der Disziplin TOPS/Watt ist die TPU um den Faktor 30 bis 80 besser. Die Forscher rechnen außerdem aus, dass man diese Werte noch steigern könnte, wenn man auf der TPU GDDR5-Speicher statt DDR3 verwenden würde. Das könnte die erzielbare Rechenleistung noch einmal verdreifachen und die Energieeffizienz auf Faktor 70 gegenüber der GPU und Faktor 200 gegenüber der CPU bringen. (bo)
