Leistungsleitfaden für Cloud TPU

Wenn Sie Probleme mit der TPU-Leistung beheben möchten, besteht der erste Schritt darin, ein Profil für Ihr Modell zu erstellen. Weitere Informationen zum Erstellen eines Leistungsprofils finden Sie unter Profil für Ihr Modell auf Cloud TPU erstellen.

Leistung des TPU-Modells

In diesem Abschnitt werden allgemeine Probleme, die die Leistung des Modells beeinträchtigen können, sowie Möglichkeiten zu deren Behebung beschrieben.

1. Das Modell ist eingabegebunden

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

2. Die Batchgröße ist 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).

   Sie nutzen den Arbeitsspeicher optimal, wenn Sie die maximale Batchgröße verwenden, 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. Optimierung der Speicherverwaltung

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

• Beschreibung: Mit TPU_PREMAPPED_BUFFER_SIZE wird die Größe des Host-Zwischenspeichers (in Byte) festgelegt, 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 * 1.024 Byte = 4.096 = 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: Eine Erhöhung dieser Größe kann die Datenübertragungsleistung zwischen dem Host und dem TPU-Gerät verbessern, insbesondere bei Arbeitslasten mit großen Tensoren oder häufiger Kommunikation zwischen Host und Gerät. Allerdings erhöht sich dadurch auch die Größe des angepinnten Host-Speichers, wodurch weniger Speicher für andere Prozesse zur Verfügung steht.

  Zwischenspeichergröße

  Wenn der vorab zugeordnete Zwischenspeicherbereich nicht groß genug ist, um während der Programmlaufzeit Speicher zuzuweisen, ist die Arbeitslast nicht möglich und gibt einen RESOURCE_EXHAUSTED-Fehler ähnlich dem folgenden zurück:

  „Die Zuweisung des Zwischenspeichers aus dem vorab zugeordneten Bereich ist fehlgeschlagen mit: RESOURCE_EXHAUSTED. Es wurde versucht, allocation_size zuzuweisen. Das war nicht möglich. Es sind available_size frei.“

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

  Sehen Sie sich zur Diagnose die TPU-Laufzeitprotokolle an. Diese Protokolle enthalten detaillierte Informationen zu den ausgeführten Vorgängen, einschließlich der Vorabzuordnung von Zwischenspeichern. Sie finden die Protokolle unter /tmp/tpu_logs/tpu_driver.INFO oder können sie direkt auf der Konsole ausgeben, indem Sie die Umgebungsvariable TPU_STDERR_LOG_LEVEL=0 festlegen. Diese Einstellung generiert Ausgabe ähnlich wie die folgende:

   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: Ist der vorab zugeordnete Zwischenspeicher zu klein oder zu groß, können Sie seine Größ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.
  

  Sie können zum Beispiel mit

   export TPU_PREMAPPED_BUFFER_SIZE=4294967296
  

  die Zwischenspeichergröß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 einen Engpass darstellt. Berücksichtigen Sie die Hostspeichernutzung und die Modellleistung, um das optimale 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-Codebasis von TensorFlow, kann aber auch für PyTorch- und JAX-Modelle verwendet werden. Modelle für Cloud TPU werden in eine XLA-Grafik übertragen und dann in eine ausführbare TPU-Datei kompiliert. Weitere Informationen zu XLA finden Sie unter XLA: Compiler für maschinelles Lernen optimieren.

Padding

Um den TPU-Speicher effizient zu nutzen, strukturieren Sie Ihre Daten so, dass sie in 128 x 8 Blöcke unterteilt werden können. Wenn die Daten für eine Matrixberechnung keinen vollständigen 128 × 8-Bock 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.

Das Padding wird bei Bedarf automatisch vom XLA-Compiler durchgeführt. Den Umfang können Sie jedoch über das Tool Memory Viewer festlegen. Sie können Padding vermeiden, indem Sie Tensor-Dimensionen auswählen, die für TPU gut geeignet sind.

Tensor-Dimensionen

Um die maximale Anzahl an FLOPs zu erreichen, sollten die Dimensionen der Matrixmultiplikation die MXU-Größe für die verwendete TPU-Version übersteigen. 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 der im TPU-HBM-Speicher gespeicherten Tensoren auf, um Berechnungen effizienter durchzuführen. Dieses Padding erfolgt transparent auf Hardwareebene und hat keinen Einfluss 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 Rechenleistung maximiert und das Padding minimiert wird. Wenn der Speicheraufwand minimiert und die Rechenleistung 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-Dimensionen sollten ein Vielfaches von 128 sein.

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

Eine Batchgröße von 1024 und Feature-Dimensionen, die ein Vielfaches von 128 sind, führen zu der besten Effizienz. Dies ist jedoch möglicherweise nicht für alle Modelle möglich.

Hinweis: Feature-Dimension bezieht sich auf die versteckte Größe einer vollständig verbundenen Ebene oder die Anzahl der Ausgabekanäle in einer Faltung. Nicht alle Ebenen können dieser Regel entsprechen, insbesondere nicht die erste und letzte Ebene des Netzwerks. Das ist in Ordnung, die meisten Modelle benötigen eine gewisse Menge an Padding.

Fusion

Fusion ist eine allgemeine Technik, die der XLA-Compiler zur Optimierung von Programmen verwendet. Eine fusionierte Operation ist die Kombination mehrerer Teiloperationen, die zusammen 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 „tmp“ keinen Speicherplatz 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 zwar oft mit ihren Nutzern fusioniert werden, aber das Erzwingen einer Übertragung kann zu schlechter Leistung 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)`