Leistungsleitfaden für Cloud TPU

Der erste Schritt bei der Fehlerbehebung der TPU-Leistung ist die Profilerstellung für Ihr Modell. Weitere Informationen zum Erstellen eines Leistungsprofils finden Sie unter Modell auf 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. Modell ist eingabebeschränkt

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

2. Batchgröße aufgrund von Sharding (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 8 × 128-Vektorregister für die Verarbeitung von Matrixmultiplikationen. Im Allgemeinen sollte die Batchgröße gleichmäßig durch 8 oder 128 teilbar sein.

3. Speicherverwaltung optimieren

   Mit den Umgebungsvariablen TPU_PREMAPPED_BUFFER_SIZE können Sie das Laufzeitverhalten auf niedriger Ebene optimieren.

• Beschreibung:TPU_PREMAPPED_BUFFER_SIZE legt die Größe des Host-Arbeitsspeicherpuffers (in Byte) fest, der für die Verwendung durch die TPU-Laufzeit für Datenübertragungen (z. B. DMA) vorab zugeordnet und angepinnt wird. Der Standardwert ist 4.294.967.296 Byte. Der Wert muss ein Vielfaches von 2^12 (4 KB = 4 * 1024 Byte = 4096 = 2^12) sein.

  Die folgenden Beispiele sind gültige TPU_PRE_MAPPED_BUFFER_SIZE-Werte.

      17179869184 = 2^34 = 2^22 * 2^12 (2^22 4KB pages will be premapped).
      40000000000 = 5^10 * 2^12 = (5^10 4KB pages will be premapped).
  

• Auswirkung:Wenn Sie diese Größe erhöhen, kann die Leistung der Datenübertragung zwischen dem Host und dem TPU-Gerät verbessert werden, insbesondere bei Arbeitslasten mit großen Tensoren oder häufiger Host-Geräte-Kommunikation. Dadurch wird jedoch auch die Menge des angepinnten Hostspeichers erhöht, wodurch der für andere Prozesse verfügbare Speicher reduziert wird.

  Puffergröße

  Wenn der vorab zugeordnete Pufferbereich nicht groß genug ist, um während der Programmlaufzeit Speicher zuzuweisen, schlägt die Arbeitslast fehl und es wird ein RESOURCE_EXHAUSTED-Fehler ähnlich dem folgenden zurückgegeben:

  „Das Zuweisen des Puffers aus der vorab zugeordneten Region ist fehlgeschlagen: RESOURCE_EXHAUSTED. Es wird versucht, allocation_size zuzuweisen. Das war nicht möglich. Es gibt available_size kostenlose.“

  Wenn der Puffer zu groß ist, kann die TPU-Initialisierung viel länger dauern (möglicherweise mehr als 15 Sekunden), sodass es so aussieht, als ob die TPU hängen geblieben ist.

  Sehen Sie sich zur Diagnose die TPU-Laufzeitprotokolle an. Diese Logs enthalten Details zu den ausgeführten Vorgängen, einschließlich der Vorabzuordnung von Puffern. Sie finden die Logs unter /tmp/tpu_logs/tpu_driver.INFO oder können sie direkt in der Konsole ausgeben lassen, indem Sie die Umgebungsvariable TPU_STDERR_LOG_LEVEL=0 festlegen. Diese Einstellung generiert eine Ausgabe, die etwa so aussieht:

   I0604 12:45:24.926233   62136 tpu_hal.cc:214] Starting premapped memory manager initialization...
   I0604 12:45:29.411218   62136 system.cc:1059] tpu::System initialized, current host id: 0, logical device ids: 0
   I0604 12:45:29.411244   61600 tfrt_tpu_system_state.cc:216] CreateTpuSystemState: TPU initialization is successful and it took 5.583190661s
   I0604 12:45:29.411267   61600 tfrt_tpu_system_state.cc:220] CreateTpuSystemState: using TPU host premapped buffer of size: 4294967296
   ```
  
  This output will tell you how long it took to initialize the TPU and
  the size of the premapped buffer.
  

• Verwendung:Wenn der vorab zugeordnete Puffer zu klein oder zu groß ist, können Sie die Puffergröße manuell mit den folgenden Umgebungsvariablen festlegen.

  TPU_PREMAPPED_BUFFER_SIZE: Sets the total size (in bytes) of the
  pre-mapped buffer region.
  TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES: Sets the maximum size of
  a single buffer that can be allocated from the pre-mapped region.
  

  Beispiele:

   export TPU_PREMAPPED_BUFFER_SIZE=4294967296
  

  die Puffergröße festlegen und:

   export TPU_PREMAPPED_BUFFER_TRANSFER_THRESHOLD_BYTES
   ```
   to enable it.
  
   This export sets the size to the default.
  

• Hinweis:Passen Sie den Wert von TPU_PREMAPPED_BUFFER_SIZE an, wenn Sie vermuten, dass die Datenübertragung zwischen Host und Gerät ein Engpass ist. Überwachen Sie die Hostspeichernutzung und die Modellleistung, um ein optimales Gleichgewicht zu finden. Der Standardwert ist in der Regel für die meisten Anwendungsfälle ausreichend.

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 Standard-TensorFlow-Codebasis, kann aber auch für PyTorch- und JAX-Modelle verwendet werden. Modelle für Cloud TPU werden in ein XLA-Diagramm übersetzt, das XLA dann in eine ausführbare TPU-Datei kompiliert. Weitere Informationen zu XLA finden Sie unter XLA: Optimizing Compiler for Machine Learning.

Padding

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

1. Aufgefüllte Tensoren lasten 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.

Tensor-Dimensionen

Um die maximale Anzahl an FLOPS zu erreichen, sollten die Dimensionen der Matrixmultiplikation größer als die MXU-Größe für die verwendete TPU-Version sein. Die MXU-Größe beträgt 256 × 256 für v6e und 128 × 128 für Versionen vor v6e. Weitere Informationen finden Sie unter Cloud TPU-Systemarchitektur.

Batchgröße

Der XLA-Compiler rundet die Größen von Tensoren, die im TPU-HBM-Speicher gespeichert sind, auf, um Berechnungen effizienter durchzuführen. 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 so an, dass die Recheneffizienz maximiert und das Padding minimiert wird. 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 Feature-Dimensionsgrößen sollten ein Vielfaches von 128 sein.

2. Die Gesamt-Batchgröße sollte ein Vielfaches von 1.024 (128 pro TPU-Kern) sein. Die Größen der Feature-Dimensionen sollten 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.

Hinweis :Die Feature-Dimension bezieht sich auf die verborgene Größe einer vollständig verbundenen Ebene oder die Anzahl der Ausgabekanäle in einer Faltung. Nicht alle Ebenen können dieser Regel entsprechen. Dies gilt insbesondere für die erste und die letzte Ebene des Netzwerks. Das ist in Ordnung, die meisten Modelle erfordern ein gewisses Maß an Padding.

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[i] = tmp[i] * z[i];
    }

Bei der Fusion treten die Arrayzugriffe gleichzeitig auf:

    for (i = 0; i < element_count; i++) {
      result[i] = (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, kann das Erzwingen eines Broadcasts zu einer schlechten Leistung und einer erhöhten 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)`