bfloat16 mit TensorFlow-Modellen verwenden

Studien zum maschinellen Lernen (ML) zeigen, dass viele maschinelle Lernmodelle Arithmetik mit niedrigerer Genauigkeit tolerieren können, ohne die konvergierte Genauigkeit zu beeinträchtigen. Viele Modelle erzielen mit bfloat16 die gleiche konvergente Genauigkeit wie bei der Verwendung von 32-Bit-Gleitkommazahlen. Einige Modelle zeigen sogar eine verbesserte konvergente Genauigkeit mit bfloat16.

Dieses Dokument betrifft ein Mixed Precision Training im Sinne von Speichern von Aktivierungen und Farbverläufen im Speicher mit dem Format bfloat16. (Weitere Informationen zum Mixed Precision Training finden Sie unter Mixed Precision Training). Dieses Dokument behandelt nicht die Verwendung von bfloat16 in MXU der Cloud TPU.

Die folgenden Themen gelten für ML-Modelle mit TensorFlow:

  • Beschreibung des benutzerdefinierten 16 Bit-Zentralgleitkommas von Google, bfloat16
  • Leistungsvorteile der Verwendung von bfloat16 im Arbeitsspeicher für ML-Modelle auf Hardware, die dies unterstützt, beispielsweise Cloud TPU
  • Speichern von Aktivierungen und Verläufen mit bfloat16 für ein TPU-Modell in TensorFlow

Standardmäßig speichert TensorFlow alle Variablen in 32-Bit-Gleitkomma (FP32). Die Verwendung von bfloat16 für die Aktivierungen und Farbverläufe beschleunigt die Schrittzeit des Geräts und verringert den Speicherbedarf. Siehe Ändern des Modells zur Bestimmung der Vorteile der Verwendung von bfloat16 für Aktivierungen und Gradienten in Ihrem Modell.

Das Gleitkommaformat bfloat16

Das bfloat16-Format ist [1:8:7], bestehend aus einem Vorzeichenbit, acht Exponentenbits und sieben Mantissenbits plus einem impliziten Mantissenbit. Zum Vergleich: Das Standardformat für 16-Bit-Gleitkommazahlen (fp16) ist [1:5:10]. Das fp16-Format enthält nur fünf Exponentenbits. Aufgrund dieser Eigenschaften hat bfloat16 einen größeren dynamischen Bereich als fp16 hat. Der Bereich von bfloat16 ist nützlich für Dinge wie Gradienten, die außerhalb des dynamischen Bereichs von fp16 liegen können und daher eine Verlustskalierung erfordern. bfloat16 kann solche Gradienten direkt darstellen. Sie können das Format bfloat16 außerdem zur genauen Darstellung aller Ganzzahlen [-256, 256] verwenden und dadurch ein int8 in bfloat16 ohne Genauigkeitsverlust codieren.

Die folgende Abbildung zeigt drei Gleitkommaformate:

  • fp32: IEEE-Gleitkommazahl mit einfacher Genauigkeit
  • fp16: IEEE-Gleitkomma mit halber Genauigkeit
  • bfloat16: 16 Bit-Zentralgleitkomma

Bild

Der dynamische Bereich von bfloat16 ist größer als der von fp16.

Vorteile hinsichtlich Leistung und Speicherauslastung

Cloud TPU unterstützt das Speichern von Werten wie Aktivierungen und Farbverläufen im Format bfloat16. bfloat16 reduziert die Größe der Daten im Speicher und ermöglicht, dass größere Modelle in den gleichen Speicherplatz passen. Die Verwendung von bfloat16 kann auch die Rematerialisierung reduzieren, wodurch die Schrittzeit verbessert wird.

Einige Operationen sind speicherbandbreitengebunden, d. h. die Speicherbandbreite bestimmt die für solche Operationen aufgewendete Zeit. Das Speichern von Ein- und Ausgängen von speicherbandbreitenabhängigen Operationen im bfloat16-Format reduziert die zu übertragende Datenmenge, wodurch die Geschwindigkeit der Operationen verbessert wird.

Die folgende Tabelle zeigt die Verbesserungen, die in unseren internen Experimenten festgestellt wurden.

Bild

Modell zur Verwendung von bfloat16 ändern

Aktivierungen, Verläufe und Gewichtungen werden standardmäßig in fp32 gespeichert. Sie können bfloat16 für Aktivierungen und Farbverläufe verwenden, Gewichte in fp32 belassen und dann die Leistung Ihres Modells zwischen bfloat16 und fp32 vergleichen, um die Vorteile zu ermitteln.

  1. Führen Sie das Modell in fp32 mit capture_tpu_profile aus.

  2. Verwenden Sie den Profilbetrachter in TensorBoard, um die Schrittzeit und die konvergente Genauigkeit des Modells anzuzeigen. (Weitere Informationen finden Sie unter Cloud TPU-Tools verwenden.)

  3. Konvertieren Sie die Eingabe in bfloat16 in Ihrer Eingabepipeline innerhalb des Record Parser, sodass die Konvertierung parallel und nicht am Ende von input_fn erfolgen kann. Dadurch werden alle Aktivierungen und Verläufe im Modell in bfloat16 konvertiert.

    Beispiel:

    image = tf.cast(image, tf.bfloat16)
    
  4. Erstellen Sie Ihr Netzwerk unter bfloat16 und setzen Sie die Ausgaben des Modells in float32 um.

Nachdem Sie Ihr Modell für die Verwendung von tf.bfloat16 für Aktivierungen konfiguriert haben, prüfen Sie Folgendes, um die Auswirkungen von bfloat16 auf Ihr Modell zu sehen:

  1. Führen Sie das Modell mit bfloat16 mit capture_tpu_profile aus.
  2. Verwenden Sie den Profilbetrachter in TensorBoard, um die Schrittzeit und die konvergente Genauigkeit des Modells anzuzeigen. (Weitere Informationen finden Sie unter Cloud TPU-Tools verwenden.)
  3. Vergleichen Sie die Schrittzeit für bfloat16 und fp32. Die Schrittzeit verbessert sich normalerweise für bfloat16.
  4. Vergleichen Sie die konvergente Genauigkeit für bfloat16 mit fp32. Normalerweise sind sie identisch, aber die Werte können besser oder schlechter als erwartet sein.

Wenn Sie noch nicht wissen, welchen Variationsbereich Sie von Ihrem Modell erwarten können, benötigen Sie möglicherweise mehrere Durchläufe, um die Abweichungen von Lauf zu Lauf in der konvergierten Genauigkeit zu bestimmen.

Wenn Ihr Profil zeigt, dass die Verarbeitungszeit schneller ist, die Eingabepipeline jedoch zum Engpass geworden ist, optimieren Sie die Eingabepipeline für einen noch größeren Geschwindigkeitsvorteil. Siehe Data Input Pipeline Performance für allgemeine Leitlinien zur Verbesserung der TensorFlow Pipeline Performance.

Für Training und Inferenz empfiehlt es sich, für beide die gleiche Präzision zu verwenden. Sie können mit fp32 für Aktivierungen trainieren und dann mit bfloat16 (oder umgekehrt) eine Inferenz ausführen. Wenn Sie sich für eine nicht übereinstimmende Genauigkeit entscheiden, prüfen Sie die konvergierte Genauigkeit anhand der für die Inferenz verwendeten Genauigkeit.