Diese Seite wurde von der Cloud Translation API übersetzt.
Switch to English

XLA: Optimierung des Compilers für maschinelles Lernen

XLA (Accelerated Linear Algebra) ist ein domänenspezifischer Compiler für lineare Algebra, der TensorFlow-Modelle ohne potenzielle Änderungen des Quellcodes beschleunigen kann.

Das Ergebnis sind Verbesserungen bei Geschwindigkeit und Speichernutzung: Die meisten internen Benchmarks werden nach Aktivierung von XLA ~ 1,15-mal schneller ausgeführt. Der folgende Datensatz wird auf einer einzelnen NVidia V100-GPU ausgewertet:

Einführung

Wenn ein TensorFlow-Programm ausgeführt wird, werden alle Vorgänge vom TensorFlow-Executor einzeln ausgeführt. Jede TensorFlow-Operation verfügt über eine vorkompilierte GPU-Kernel-Implementierung, an die der Executor sendet.

XLA bietet einen alternativen Modus zum Ausführen von Modellen: Es kompiliert das TensorFlow-Diagramm in eine Folge von Rechenkernen, die speziell für das jeweilige Modell generiert wurden. Da diese Kernel für das Modell eindeutig sind, können sie modellspezifische Informationen zur Optimierung nutzen. Schauen wir uns zum Beispiel eine Optimierung an, die XLA im Kontext einer einfachen TensorFlow-Berechnung durchführt:

 def model_fn(x, y, z):
  return tf.reduce_sum(x + y * z)
 

Das Diagramm wird ohne XLA ausgeführt und startet drei Kernel: einen für die Multiplikation, einen für die Addition und einen für die Reduktion. XLA kann das Diagramm jedoch so optimieren, dass das Ergebnis in einem einzelnen Kernelstart berechnet wird. Dies geschieht durch "Zusammenführen" der Addition, Multiplikation und Reduktion zu einem einzigen GPU-Kernel. Darüber hinaus schreibt diese verschmolzene Operation die von y*z und x+y*z erzeugten Zwischenwerte nicht in den Speicher; Stattdessen werden die Ergebnisse dieser Zwischenberechnungen direkt an ihre Benutzer "gestreamt", während sie vollständig in GPU-Registern gespeichert werden. Fusion ist die wichtigste Optimierung von XLA. Die Speicherbandbreite ist normalerweise die knappste Ressource für Hardwarebeschleuniger. Das Entfernen von Speicheroperationen ist daher eine der besten Möglichkeiten, um die Leistung zu verbessern.

Aktivieren Sie XLA für TensorFlow-Modelle

Auto-Clustering

Eine einfachste Möglichkeit, XLA in TensorFlow-Modellen zu verwenden, besteht darin, das automatische Clustering zu aktivieren, bei dem automatisch Cluster (verbundene Untergraphen) im TensorFlow-Diagramm gefunden werden, die mit XLA kompiliert und ausgeführt werden können. Das automatische Clustering auf der GPU kann durch Festlegen der Umgebungsvariablen TF_XLA_FLAGS werden:

 $ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program
 

Auto-Clustering ist derzeit für GPU-Workloads optimiert, kann jedoch auch auf der CPU aktiviert werden, indem zusätzlich das Flag --tf_xla_cpu_global_jit :

 $ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program
 

Ein detailliertes Verwendungsbeispiel finden Sie in der Tutorial-Anleitung zum automatischen Clustering .

Explizite Kompilierung mit tf.function

Auto-Clustering ist ein großartiges Tool, um das Modell ohne Änderungen am Code schneller zu machen. Es ist jedoch möglicherweise schwer zu verstehen, welche Änderungen vorgenommen wurden.

Die explizite Kompilierungs-API bietet eine detailliertere Steuerung für die Auswahl der zu kompilierenden Funktionen. Beispielsweise wird die folgende TensorFlow-Funktion, die das MNIST-Training durchführt, mit XLA kompiliert:

 @tf.function(experimental_compile=True)
def train_mnist(images, labels):
    images, labels = cast(images, labels)

    with tf.GradientTape() as tape:
      predicted_labels = layer(images)
      loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
          logits=predicted_labels, labels=labels
      ))
    layer_variables = layer.trainable_variables
    grads = tape.gradient(loss, layer_variables)
    optimizer.apply_gradients(zip(grads, layer_variables))
 

Die experimental_compile API verfügt über eine Semantik, die unbedingt kompiliert werden muss : Entweder wird die gesamte Funktion mit XLA kompiliert, oder es wird eine errors.InvalidArgumentError Ausnahme ausgelöst. XLA kann derzeit keine Funktionen kompilieren, bei denen Dimensionen nicht ableitbar sind . Das heißt, wenn es nicht möglich ist, die Dimensionen aller Tensoren abzuleiten , ohne die gesamte Berechnung auszuführen . Die folgende Funktion wird beispielsweise nicht kompiliert:

 @tf.function
def not_compilable(x):
  return tf.unique(x)
 

Die Formen können jedoch über die Läufe variieren:

 @tf.function(experimental_compile=True)
def recompiled_on_launch(a, b):
  return a + b

recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))
 

In der Tutorial-Spalte finden Sie ein detaillierteres Anwendungsbeispiel.

AOT-Kompilierung (Ahead-of-Time) für CPU mit tfcompile

Sie können auch ein eigenständiges tfcompile Tool verwenden, das das TensorFlow-Diagramm in ausführbaren Code konvertiert (nur für x86-64-CPU).

Überprüfen Sie kompilierte Programme

XLA bietet Introspection-Funktionen, mit denen Sie die generierten Programme überprüfen können. Verwenden Sie die Umgebungsvariable XLA_FLAGS , um die generierten Programme zu XLA_FLAGS :

 $ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program
 

Nach dem Dumping finden Sie die folgenden Dateien in /tmp/generated :

  • module_XXXX.*_optimizations.txt Generierte XLA-Programme , eines pro kompiliertem Cluster. Das Anhängen dieser beim Einreichen von XLA-Fehlerberichten ist äußerst hilfreich!

  • module_XXXX.ir-*.ll Generierte Dateien in LLVM- Zwischendarstellung mit NVPTX- Eigenschaften.

  • module_XXXX.ptx Generierte PTX- Dateien.

Sie können das Diagramm auch sichern, indem Sie die Einbettung von XLA-Clustern in das TensorFlow-Diagramm visualisieren:

 $ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"
 

Reproduzierbare Fehlerberichte

Ein Fehlerbericht ist viel einfacher zu reproduzieren, wenn er Speicherauszüge für die generierten XLA-Programme und die verwendete automatische Clustering-Einbettung enthält. Starten Sie Folgendes, um sie für ein TensorFlow-Programm zu generieren, das mit Auto-Clustering ausgeführt wird:

 $ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
  TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
  XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
    my/tensorflow/program"
 

Fügen Sie beim Einreichen von Fehlern den Inhalt des Verzeichnisses /tmp/generated (siehe oben).

Versuchen Sie nach Möglichkeit, einen Fehler in einem einzelnen XLA-Programm zu isolieren, indem Sie die replay_computation und iterativ auf generierten Programmen replay_computation .

XLA-Frontends

Neben TensorFlow können XLA-Programme generiert werden durch:

  • JAX : Zusammensetzbare Transformationen von Python + NumPy-Programmen
  • Julia : Die Julia-Sprache für wissenschaftliches Rechnen
  • PyTorch : PyTorch-Framework

Weiterführende Literatur