Einführung in Cloud TPU

Tensor Processing Units (TPUs) sind von Google speziell entwickelte, anwendungsspezifische integrierte Schaltungen zur Beschleunigung von ML-Arbeitslasten. Weitere Informationen zur 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, da sie Hardware verwenden, die für die Ausführung großer Matrixvorgänge ausgelegt ist, wie sie häufig in Algorithmen für maschinelles Lernen vorkommen. TPUs haben einen auf dem Chip integrierten Arbeitsspeicher mit hoher Bandbreite (HBM), sodass Sie größere Modelle und Batches verwenden können. TPUs können in Gruppen, sogenannten Slices, verbunden werden, die Ihre Arbeitslasten mit geringen oder gar keinen Codeänderungen skalieren.

Auf TPUs ausgeführter Code muss vom XLA-Compiler (Accelerator Linear Algebra) kompiliert werden. XLA ist ein Just-in-Time-Compiler. Er kompiliert die linearen Algebra-, Verlust- und Gradientenkomponenten des von einer ML-Framework-Anwendung ausgegebenen Graphen in TPU-Maschinencode. Das restliche Programm wird auf dem TPU-Hostcomputer ausgeführt. Der XLA-Compiler gehört zum TPU-VM-Image, das auf einem TPU-Hostcomputer ausgeführt wird.

Weitere Informationen zu Tensor Processing Units finden Sie unter How to think about TPUs.

Verwendung von TPUs

Cloud TPUs sind für bestimmte Arbeitslasten optimiert. In einigen Situationen möchten Sie für Ihre Machine-Learning-Arbeitslasten vielleicht GPUs oder CPUs auf Compute Engine-Instanzen verwenden. Normalerweise können Sie anhand der folgenden Richtlinien entscheiden, welche Hardware für Ihre Arbeitslast am besten geeignet ist.

CPUs

GPUs

  • Modelle mit einer erheblichen Anzahl benutzerdefinierter PyTorch-/JAX-Vorgänge, die zumindest teilweise auf CPUs ausgeführt werden müssen
  • Modelle mit TensorFlow-Vorgängen, die auf Cloud TPU nicht verfügbar sind (siehe Liste der verfügbaren TensorFlow-Vorgänge)
  • 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 extrem großen Einbettungen, wie sie bei anspruchsvollen Ranking- und Empfehlungs-Arbeitslasten üblich sind

Cloud TPUs sind für folgende Arbeitslasten nicht geeignet:

  • Lineare Algebraprogramme, die häufige Verzweigungen erfordern oder viele elementweise Algebravorgänge enthalten
  • Arbeitslasten, die hochpräzise Arithmetik erfordern
  • Neuronale Netzwerkarbeitslasten, die in der Haupttrainingsschleife benutzerdefinierte Vorgänge 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 die Ressourcen für die einzelnen Dienste von Google Cloudaufgefü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 Berechnungen überwiegend aus Nicht-Matrix-Vorgängen wie Addieren, Umformen oder Verknüpfen bestehen, wird wahrscheinlich keine hohe MXU-Auslastung erzielen. Es folgen einige Richtlinien, die Ihnen bei der Auswahl und Erstellung von Modellen helfen, 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. Er nutzt die Struktur der MXU-Hardware, ein systolisches Array der Größe 128 x 128, und das Design des TPU-Speichersubsystems, das Dimensionen bevorzugt, die Vielfache von 8 sind, um die Effizienz der Kachelung zu steigern.

Folglich sind bestimmte Layouts für die Kachelung besser geeignet, während andere vor der Kachelung eine Umformung erfordern. Umformungsvorgänge sind auf der Cloud TPU oft speichergebunden.

Formen

Der XLA-Compiler kompiliert einen ML-Graphen gerade rechtzeitig für den ersten Batch. Wenn nachfolgende Batches andere Formen haben, funktioniert das Modell nicht. (Das erneute Kompilieren des Graphen bei jeder Formänderung ist zu langsam.) Daher sind Modelle mit Tensoren mit dynamischen Formen für TPUs nicht gut geeignet.

Padding

Ein leistungsstarkes Cloud-TPU-Programm ist eines, bei dem dichte Berechnungen in 128x128-Blöcke gekachelt werden können. Kann eine Matrixberechnung nicht eine gesamte MXU belegen, füllt der Compiler die Tensoren mit Nullen auf. Padding hat zwei Nachteile:

  • Mit Nullen aufgefüllte Tensoren nutzen den TPU-Kern nicht vollständig aus.
  • Padding erweitert den für einen Tensor erforderlichen On-Chip-Arbeitsspeicher und kann im Extremfall zu einem Fehler aufgrund fehlenden Arbeitsspeichers führen.

Das Padding wird bei Bedarf automatisch vom XLA-Compiler durchgeführt. Den Umfang des Paddings können Sie jedoch über das Tool op_profile bestimmen. Sie können Padding vermeiden, indem Sie Tensor-Dimensionen auswählen, 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 Featuredimension 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 die Batchgröße und die Featuredimensionen Vielfache von 8 sein, damit eine hohe Leistung aus dem Speichersubsystem erzielt werden kann.

Einstieg in Cloud TPU

  1. Ein Konto in Google Cloud einrichten
  2. Die Cloud TPU API aktivieren
  3. Cloud TPU Zugriff auf Ihre Cloud Storage-Buckets gewähren
  4. Eine einfache Berechnung auf einer TPU ausführen
  5. Ein Referenzmodell auf einer TPU trainieren
  6. Ihr Modell analysieren

Hilfe anfordern

Wenn Sie Hilfe benötigen, wenden Sie sich an den Cloud TPU-Support. Wenn Sie ein aktives Projekt in Google Cloud haben, halten Sie die folgenden Informationen bereit:

  • Ihre Projekt-ID in Google Cloud
  • Der Name Ihrer TPU, falls 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: