Cloud TPU-Leistungsleitfaden

Der erste Schritt bei der Fehlerbehebung bei der TPU-Leistung besteht darin, ein Profil für Ihr Modell zu erstellen. Weitere Informationen zum Erfassen eines Leistungsprofils finden Sie unter Profilerstellung für Ihr Modell in Cloud TPU.

Leistung des TPU-Modells

In diesem Abschnitt werden allgemeine Probleme beschrieben, die die Modellleistung beeinträchtigen können, und wie Sie diese beheben.

  1. Modell ist eingabegebunden

    TPUs führen Berechnungen sehr schnell durch. Damit die TPU nicht inaktiv ist, muss ein stetiger Datenstrom auf die TPU geladen werden. Wie dies funktioniert, hängt davon ab, wie Sie Ihr Dataset laden und vorverarbeiten. Sie können beispielsweise Datendateien mit tf.data.TFRecordset() und dem Parameter num_parallel_reads parallel lesen.

  2. Batchgröße aufgrund von Fragmentierung (Aufteilung von Batches auf Kerne) zu klein

    Die TPU-Laufzeit teilt einen Batch auf alle 8 Kerne eines TPU-Geräts (z. B. v2-8 oder v3-8) auf. Wenn Sie eine globale Batchgröße von 128 angeben, erhält jeder Kern eine Batchgröße von 16 (128 / 8).

    Verwenden Sie für eine optimale Arbeitsspeichernutzung die größte Batchgröße, die in den TPU-Arbeitsspeicher passt. Jeder TPU-Kern verwendet zweidimensionale 8 × 128-Vektorregister zur Verarbeitung von Matrixmultiplikationen. Im Allgemeinen sollte die Batchgröße gleichmäßig durch 8 oder 128 teilbar sein.

XLA-Compiler-Optimierungen

XLA ist ein Compiler für maschinelles Lernen, der Binärprogramme für TPUs, CPUs, GPUs und andere Plattformen erstellen kann. XLA ist zwar Teil der TensorFlow-Standardcodebasis, kann aber auch für PyTorch- und JAX-Modelle verwendet werden. Modelle für Cloud TPU werden in eine XLA-Grafik übersetzt, die XLA dann in eine ausführbare TPU kompiliert. Weitere Informationen zu XLA finden Sie unter XLA: Optimierung des Compilers für maschinelles Lernen.

Padding

Strukturieren Sie Ihre Daten so, dass sie in Blöcken von 128 x 8 gekachelt werden können, um den TPU-Arbeitsspeicher effizient nutzen zu können. Wenn die Daten für eine Matrixberechnung nicht einen ganzen Block von 128 x 8 füllen, füllt der XLA-Compiler Tensoren auf. Das Auffüllen hat zwei Nachteile:

  1. Gepolsterte Tensoren nutzen den TPU-Kern zu wenig aus.
  2. Das Padding erhöht den für einen Tensor erforderlichen On-Chip-Speicher und kann zu einem Fehler aufgrund fehlenden Arbeitsspeichers führen.

Während das Padding bei Bedarf automatisch vom XLA-Compiler ausgeführt wird, können Sie den Umfang des Paddings mit dem Arbeitsspeicher-Viewer-Tool bestimmen. Sie können ein Padding vermeiden, indem Sie Tensor-Dimensionen auswählen, die für TPU gut geeignet sind.

Tensor-Dimensionen

Der XLA-Compiler rundet die Größen der im TPU-HBM-Speicher gespeicherten Tensoren auf, um Berechnungen effizienter auszuführen. Das Padding erfolgt transparent auf Hardwareebene und hat keine Auswirkungen auf die Ergebnisse. In bestimmten Fällen kann Padding jedoch zu einer deutlich erhöhten Speicherauslastung und Ausführungszeit führen.

Die TPU-Laufzeit legt Tensoren im Speicher an, um die Recheneffizienz zu maximieren und das Padding zu minimieren. Um den Speicheraufwand zu minimieren und die Recheneffizienz zu maximieren, muss eine der folgenden Bedingungen erfüllt sein:

  1. Die Gesamt-Batchgröße sollte ein Vielfaches von 64 (8 pro TPU-Kern) und die Größen der Feature-Dimensionen ein Vielfaches von 128 sein.

  2. Die Gesamt-Batchgröße sollte ein Vielfaches von 1024 (128 pro TPU-Kern) sein und die Größen der Feature-Dimensionen sollte ein Vielfaches von 8 sein.

Die Verwendung einer Batchgröße von 1.024 und von Featuredimensionen, die ein Vielfaches von 128 sind, ermöglicht eine optimale Effizienz, obwohl dies möglicherweise nicht für alle Modelle möglich ist.

Fusion

Fusion ist eine allgemeine vom XLA-Compiler verwendete Technik zur Optimierung von Programmen. Eine fusionierte Operation ist die Kombination mehrerer Bestandteile, die in Kombination ausgeführt werden sollen.

Betrachten Sie beispielsweise die folgende Reihe von Operationen:

    tmp = tf.add(x, y)
    result = tf.multiply(tmp, z)

Dieser Code entspricht in etwa dem folgenden Pseudocode:

    for (i = 0; i < element_count; i++) {
      tmp[i] = x[i] + y[i];
    }

    for (i = 0; i < element_count; i++) {
      result = tmp[i] * z[i];
    }

Bei der Fusion treten die Arrayzugriffe gleichzeitig auf:

    for (i = 0; i < element_count; i++) {
      result = (x[i] + y[i]) * z[i];
    }

In diesem Beispiel wird die Anzahl der Speicherumläufe reduziert und XLA muss keinen Platz für „tmp“ zuweisen.

Fusion ist eine entscheidende Optimierung und kommt der Cloud TPU auf verschiedene Arten zugute:

  • Es reduziert Arbeitsspeicherübertragungen, da Zwischenergebnisse nicht im Hauptarbeitsspeicher gespeichert werden müssen, was langsam ist.
  • Es ermöglicht eine bessere Nutzung von Hardwarekomponenten, die sonst ungenutzt wären.
  • Es kann die Arbeitsspeichernutzung eines Modells verringern, da weniger Puffer gleichzeitig aktiv sein müssen.

Broadcasting

Broadcasting tritt implizit auf, wenn zwei Tensoren mit verschiedenen, aber kompatiblen Formen kombiniert werden.

Für tf.add(vector, matrix) muss der Vektor beispielsweise an die Form der Matrix gesendet werden. Das Ergebnis der Operation hat die gleiche Form wie die Matrix. Weitere Informationen finden Sie in der Anleitung zum Übertragen von Arrays.

Obwohl Broadcasts häufig mit ihren Nutzern zusammengeführt werden können, kann das Erzwingen einer Übertragung zu schlechter Leistung und erhöhter Speichernutzung führen.

Im folgenden Beispiel kann die implizit im ergänzenden Vektor oder der ergänzenden Matrix enthaltene Übertragung nicht mit argmax zusammengeführt werden, was zu einer materialisierten Übertragung führt:

`tf.argmax(tf.add(vector, zero_matrix), axis=0)`