Leistungsleitfaden für Cloud TPU

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

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 durch. Um sicherzustellen, dass die TPU nicht inaktiv ist, ist es wichtig, dass ein kontinuierlicher Datenstrom TPU 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 (für Beispiel v2-8 oder v3-8). Wenn Sie eine globale Batchgröße von 128 angeben, erhält jeder Kern Batchgröße von 16 (128 / 8).

    Verwenden Sie für eine optimale Arbeitsspeichernutzung die größte Batchgröße, die in eine TPU passt zu speichern. Jeder TPU-Kern verwendet zweidimensionale Vektorregister (8 × 128) Matrixmultiplikationen zu verarbeiten. Im Allgemeinen sollte die Batchgröße durch 8 oder 128 teilbar sein.

XLA-Compiler-Optimierungen

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 zwar Teil der standardmäßigen TensorFlow-Codebasis, kann aber auch auf PyTorch- und JAX-Modelle Modelle für Cloud TPU in ein XLA-Diagramm übersetzt, das XLA dann in eine TPU kompiliert. ausführbar sein. Weitere Informationen zu XLA finden Sie unter XLA: Optimierter Compiler für maschinelles Lernen.

Padding

Strukturieren Sie Ihre Daten so, dass sie in Kacheln gekachelt werden können, um den TPU-Arbeitsspeicher effizient zu nutzen. 128 x 8 Blöcke. Wenn die Daten für eine Matrixberechnung nicht einen ganzen 128-x-8-Chunk, der XLA-Compiler füllt Tensoren auf. Das 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 Auffüllen bei Bedarf automatisch vom XLA-Compiler durchgeführt wird, können Sie den mit dem Memory Viewer durchgeführten Padding-Wert bestimmen. Padding kann durch Auswahl von Tensordimensionen vermieden werden, die für TPUs gut geeignet sind.

Tensorabmessungen

Der XLA-Compiler rundet die Größe der im TPU-HBM-Speicher gespeicherten Tensoren auf Berechnungen effizienter durchführen. Dieser Abstand erfolgt transparent an 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 bietet Tensoren im Speicher, um die Recheneffizienz zu maximieren und den Abstand 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 (8 pro TPU-Kern) sein und Feature Dimensionsgrößen ein Vielfaches von 128 sein.

  2. Die Gesamt-Batchgröße sollte ein Vielfaches von 1024 (128 pro TPU-Kern) sein. Die Größen der Feature-Dimensionen müssen 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. A kombinierte Operation ist die Kombination mehrerer konstituierender Operationen, in Kombination auszuführen.

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. XLA Sie müssen Speicherplatz 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.

Übertragungen können oft mit ihren Nutzern zusammengeführt werden. kann zu Leistungseinbußen und erhöhter Arbeitsspeichernutzung 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)`