Einführung in Cloud TPU

Tensor Processing Units (TPUs) sind von Google speziell entwickelte anwendungsspezifische integrierte Schaltungen (Application-Specific Integrated Circuits, ASICs), die verwendet werden, um die beim maschinellen Lernen entstehenden Arbeitslasten zu beschleunigen. Weitere Informationen zu TPU-Hardware finden Sie unter TPU-Architektur. Cloud TPU ist ein Webdienst, der TPUs als skalierbare Rechenressourcen auf Google Cloudzur Verfügung stellt.

TPUs trainieren Ihre Modelle effizienter mit Hardware, die für die Ausführung großer Matrixvorgänge entwickelt wurde, die oft in Algorithmen für maschinelles Lernen verwendet werden. TPUs haben einen On-Chip-Arbeitsspeicher mit hoher Bandbreite (HBM), sodass Sie größere Modelle und Batchgrößen verwenden können. TPUs können in Gruppen verbunden werden, die als Slices bezeichnet werden und die Ihnen das Skalieren Ihrer Arbeitslasten mit wenig oder ohne Codeänderungen ermöglichen.

Code, der auf TPUs ausgeführt wird, muss vom XLA-Compiler (Accelerator Linear Algebra) kompiliert werden. XLA ist ein Just-in-Time-Compiler, der den von einer ML-Framework-Anwendung gesendeten Graphen verwendet und die linearen Algebra-, Verlust- und Gradientenkomponenten des Graphen in TPU-Maschinencode kompiliert. Der Rest des Programms wird auf dem TPU-Hostcomputer ausgeführt. Der XLA-Compiler ist Teil des TPU-VM-Images, das auf einem TPU-Hostcomputer ausgeführt wird.

Verwendung von TPUs

Cloud TPUs sind für bestimmte Arbeitslasten optimiert. In einigen Situationen möchten Sie möglicherweise GPUs oder CPUs auf Compute Engine-Instanzen verwenden, um Ihre Arbeitslasten für maschinelles Lernen auszuführen. Im Allgemeinen können Sie basierend auf den folgenden Richtlinien entscheiden, welche Hardware für Ihre Arbeitslast am besten geeignet ist:

CPUs

GPUs

  • Modelle mit einer erheblichen Anzahl von benutzerdefinierten PyTorch-/JAX-Vorgängen, die mindestens teilweise auf CPUs ausgeführt werden müssen
  • Modelle mit TensorFlow-Operationen, die nicht auf Cloud TPU verfügbar sind (siehe die Liste der verfügbaren TensorFlow-Operationen)
  • Mittlere bis große Modelle mit größeren effektiven Batchgrößen

TPUs

  • Modelle, die von Matrixberechnungen dominiert werden
  • Modelle ohne benutzerdefinierte PyTorch-/JAX-Vorgänge innerhalb der Haupttrainingsschleife
  • Modelle, die Wochen oder Monate lang trainiert werden
  • Große Modelle mit großen, effektiven Batchgrößen
  • Modelle mit sehr großen Einbettungen, die bei erweiterten Ranking- und Empfehlungslasten häufig vorkommen

Cloud TPUs sind für folgende Arbeitslasten nicht geeignet:

  • Lineare Algebraprogramme, die häufige Verzweigungen erfordern oder viele elementweise algebraische Operationen enthalten
  • Arbeitslasten, die eine hochpräzise Arithmetik erfordern
  • Neuronale Netzwerkarbeitslasten, die benutzerdefinierte Vorgänge in der Haupttrainingsschleife enthalten

TPUs in Google Cloud

Sie können TPUs über Cloud TPU-VMs, Google Kubernetes Engine und Vertex AI verwenden. In der folgenden Tabelle sind Ressourcen für jeden Google Cloud-Dienst aufgeführt.

Google Cloud -Dienst Ressourcen
Cloud TPU Erste Schritte mit Cloud TPU-VMs
Google Kubernetes Engine

TPUs in GKE

Ray in GKE mit TPUs ausführen
Vertex AI

Training in Vertex AI mit TPUs

TPUs für Onlinevorhersagen in Vertex AI verwenden

Best Practices für die Modellentwicklung

Ein Programm, dessen Berechnung von Nicht-Matrixoperationen wie Hinzufügen, Umformen oder Verketten dominiert wird, erreicht wahrscheinlich keine hohe MXU-Auslastung. Nachfolgend finden Sie einige Richtlinien zur Auswahl und zum Erstellen von Modellen, die für Cloud TPU geeignet sind.

Layout

Der XLA-Compiler führt Codetransformationen durch, einschließlich einer Matrixmultiplikation in kleinere Blöcke, um Berechnungen auf der Matrixeinheit (MXU) effizient auszuführen. Die Struktur der MXU-Hardware, ein systolisches Array der Größe 128 x 128, und das Design des TPU-Speichersubsystems, das Dimensionen bevorzugt, die ein Vielfaches von 8 sind, werden vom XLA-Compiler für die Kacheleffizienz verwendet. Folglich sind bestimmte Layouts besser für das Kacheln geeignet, während andere eine Umformung erfordern, bevor sie gekachelt werden können. Umformungsoperationen sind auf der Cloud TPU häufig speichergebunden.

Formen

Der XLA-Compiler kompiliert zur genau richtigen Zeit einen ML-Graph für den ersten Batch. Wenn nachfolgende Stapel unterschiedliche Formen haben, funktioniert das Modell nicht. (Wenn der Graph bei jeder Änderung der Form neu kompiliert wird, ist dies zu langsam.) Daher ist jedes Modell, das Tensoren mit dynamischen Formen aufweist, nicht gut für TPUs geeignet.

Padding

Ein Cloud TPU-Programm ist leistungsstark, wenn die Dichteberechnung in Blöcke von 128 x 128 gekachelt werden kann. Wenn eine Matrixberechnung keine vollständige MXU belegen kann, füllt der Compiler Tensoren mit Nullen auf. Padding hat zwei Nachteile:

  • Tensoren, die mit Nullen aufgefüllt sind, lasten den TPU-Kern nicht voll aus.
  • Das Padding erhöht den für einen Tensor erforderlichen On-Chip-Arbeitsspeicher und kann im Extremfall zu einem Fehler aufgrund fehlenden Arbeitsspeichers führen.

Während das Padding bei Bedarf automatisch vom XLA-Compiler durchgeführt wird, kann der Umfang des Paddings mit dem Werkzeug op_profile ermittelt werden. Padding kann durch Auswahl von Tensordimensionen vermieden werden, die für TPUs gut geeignet sind.

Dimensionen

Die Wahl geeigneter Tensordimensionen hilft bei der Maximierung der Leistung der TPU-Hardware, insbesondere der MXU. Der XLA-Compiler versucht entweder die Batchgröße oder eine Feature-Dimension zu verwenden, um die MXU maximal zu nutzen. Daher muss eine von diesen ein Vielfaches von 128 sein. Andernfalls füllt der Compiler eine von ihnen auf 128 auf. Idealerweise sollten sowohl die Batchgröße als auch die Feature-Dimensionen ein Vielfaches von 8 sein, wodurch eine hohe Leistung aus dem Speichersubsystem extrahiert werden kann.

Einstieg in Cloud TPU

Hilfe anfordern

Wenden Sie sich an den Cloud TPU-Support. Wenn Sie ein aktives Google Cloud Projekt haben, müssen Sie folgende Informationen angeben:

  • Ihre Google Cloud Projekt-ID
  • Der Name Ihrer TPU, sofern vorhanden
  • Weitere Informationen, die Sie angeben möchten

Nächste Schritte

Möchten Sie mehr über Cloud TPU erfahren? Die folgenden Ressourcen können Ihnen helfen: