Leitfaden zur Leistung von Cloud TPU

Der erste Schritt bei der Fehlerbehebung bei der TPU-Leistung besteht darin, Ihr Modell zu profilieren. Weitere Informationen zum Erfassen eines Leistungsprofils finden Sie unter Modell in Cloud TPU profilerstellen.

TPU-Modellleistung

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

  1. Das Modell ist eingabegebunden

    TPUs führen Berechnungen sehr schnell aus. Damit die TPU nicht inaktiv ist, muss ein kontinuierlicher Datenstream auf die TPU geladen werden. Wie das funktioniert, hängt davon ab, wie Sie Ihren Datensatz 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 auf (z. B. v2-8 oder v3-8). 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 Speichernutzung die größte Batchgröße, die in den TPU-Speicher passt. Jeder TPU-Kern verwendet zweidimensionale Vektorregister mit 8 × 128 Elementen für die Verarbeitung von Matrixmultiplikationen. Im Allgemeinen sollte die Batchgröße durch 8 oder 128 teilbar sein.

XLA-Compileroptimierungen

XLA ist ein Compiler für maschinelles Lernen, mit dem Binärdateien für TPUs, CPUs, GPUs und andere Plattformen erstellt werden können. XLA ist Teil der Standard-TensorFlow-Codebasis, 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-Datei kompiliert. Weitere Informationen zu XLA finden Sie unter XLA: Optimierter Compiler für maschinelles Lernen.

Padding

Um den TPU-Speicher effizient zu nutzen, strukturieren Sie Ihre Daten so, dass sie in Blöcke von 128 x 8 gekachelt werden können. Wenn die Daten für eine Matrixberechnung nicht einen ganzen 128 × 8-Chunk füllen, füllt der XLA-Compiler Tensoren auf. Padding hat zwei Nachteile:

  1. Aufgefüllte Tensoren nutzen den TPU-Kern nicht voll aus.
  2. Das Padding erhöht den für einen Tensor erforderlichen On-Chip-Arbeitsspeicher und kann 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 Tool „Memory Viewer“ ermittelt werden. Padding kann durch Auswahl von Tensordimensionen vermieden werden, die für TPUs gut geeignet sind.

Tensordimensionen

Der XLA-Compiler rundet die Größen von Tensoren ab, die im TPU-HBM-Speicher gespeichert sind, damit Berechnungen effizienter durchgeführt werden. Das Padding erfolgt auf transparente Weise auf der 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 aus, um die Rechenleistung zu maximieren und das Padding zu minimieren. Wenn der Speicheraufwand minimiert und die Recheneffizienz maximiert werden soll, muss eine der folgenden Bedingungen zutreffen:

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

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

Die Verwendung einer Batchgröße von 1.024 und von Feature-Dimensionen, die ein Vielfaches von 128 sind, ermöglicht eine optimale Effizienz, obwohl dies unter Umständen nicht für alle Modelle möglich ist.

Fusion

Fusion ist eine allgemeine vom XLA-Compiler verwendete Technik zum Optimieren von Programmen. Eine fusionierte Operation ist die Kombination aus mehreren konstituierenden Operationen, 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 für 'tmp' keinen Speicherplatz reservieren.

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.

Während Broadcasts oft mit ihren Konsumenten zusammengeführt werden können, führt das erzwungene Materialisieren einer Übertragung möglicherweise zu einer schlechten Leistung und einer erhöhten Speichernutzung.

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)`