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: