XLA: Compiler für maschinelles Lernen optimieren

OpenXLA ist ein domainspezifischer Compiler für lineare Algebra, der TensorFlow-Modelle potenziell ohne Änderungen am Quellcode beschleunigen kann.

Einleitung

Wenn ein TensorFlow-Programm ausgeführt wird, werden alle Vorgänge vom TensorFlow-Executor einzeln ausgeführt. Jeder TensorFlow-Vorgang hat eine vorkompilierte GPU-Kernelimplementierung, an die der Executor sendet.

XLA bietet einen alternativen Modus zum Ausführen von Modellen: Er kompiliert das TensorFlow-Diagramm in eine Folge von Rechen-Kernels, die speziell für das jeweilige Modell generiert wurden. Da diese Kernel nur für das Modell verwendet werden, können sie modellspezifische Informationen zur Optimierung ausnutzen. Betrachten wir beispielsweise eine Optimierung, die XLA im Kontext einer einfachen TensorFlow-Berechnung durchführt:

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

Wenn Sie ohne XLA ausgeführt werden, startet das Diagramm drei Kernel: einen für die Multiplikation, einen für die Addition und einen für die Reduzierung. XLA kann die Grafik jedoch so optimieren, dass das Ergebnis zu einem einzelnen Kernel-Start berechnet wird. Dies geschieht durch die Zusammenführung, Multiplikation und Reduktion zu einem einzigen GPU-Kernel. Außerdem werden bei diesem fusionierten Vorgang die von y*z und x+y*z generierten Zwischenwerte nicht in den Speicher geschrieben. Stattdessen werden die Ergebnisse dieser Zwischenberechnungen direkt an die Nutzer „streamt“, während sie vollständig in GPU-Registern gespeichert sind. Die Fusion ist die wichtigste Optimierung von XLA. Die Arbeitsspeicherbandbreite ist in der Regel die knappste Ressource für Hardwarebeschleuniger, daher ist das Entfernen von Arbeitsspeichervorgängen eine der besten Möglichkeiten zur Leistungsverbesserung.

XLA für TensorFlow-Modelle aktivieren

Explizite Kompilierung mit tf.function(jit_compile=True)

Die Explizite Compilation API bietet eine differenzierte Kontrolle über die Auswahl der Funktionen, die kompiliert werden sollen. Beispielsweise wird die folgende TensorFlow-Funktion, die das MNIST-Training ausführt, mit XLA kompiliert:

@tf.function(jit_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 jit_compile API hat eine Must-Compile-Semantik: Entweder wird die gesamte Funktion mit XLA kompiliert oder eine errors.InvalidArgumentError-Ausnahme wird ausgelöst. XLA kann derzeit keine Funktionen kompilieren, bei denen Dimensionen nicht ableitungsfähig 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)

Formen können jedoch während der Ausführungen variieren:

@tf.function(jit_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]))

Ein ausführlicheres Anwendungsbeispiel finden Sie im Tutorial-Collab und in einem Tutorial-Video zur Verwendung von jit_compile=True.

Verwendung mit Keras

Bei Keras-Modellen kann jit_compile=True als Argument für model.compile festgelegt werden:

model.compile(optimizer="adam", jit_compile=True)

Nutzung mit verteilter Strategie

XLA:GPU kann mit der verteilten TF-Strategie (MirroredStrategy oder MultiWorkerMirroredStrategy) verwendet werden, indem die Schrittfunktion mit jit_compile=True annotiert wird:

@tf.function(jit_compile=True)
def step_fn():
  t = tf.ones(shape=[100], dtype=tf.float32)
  ctx = tf.distribute.get_replica_context()
  return ctx.all_reduce(tf.distribute.ReduceOp.SUM, t)

@tf.function
def run_fn():
  return strategy.run(step_fn)

Automatisches Clustering

Eine einfache Möglichkeit, XLA ohne Änderungen in TensorFlow-Modellen zu verwenden, ist die Aktivierung der automatischen Clustererstellung. Dabei werden automatisch Cluster (verbundene Untergrafiken) in den TensorFlow-Funktionen gefunden, die mit XLA kompiliert und ausgeführt werden können. Sie können das automatische Clustering auf der GPU aktivieren, indem Sie die Umgebungsvariable TF_XLA_FLAGS festlegen:

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

Das automatische Clustering ist derzeit für GPU-Arbeitslasten optimiert, kann aber auch auf der CPU aktiviert werden, indem Sie zusätzlich das Flag --tf_xla_cpu_global_jit verwenden:

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

Ein detailliertes Anwendungsbeispiel finden Sie in der Anleitung zum automatischen Clustering in Colab.

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

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

Kompilierte Programme prüfen

XLA bietet Möglichkeiten zur Selbstprüfung, mit denen Sie die generierten Programme untersuchen können. Verwenden Sie zum Sichern der generierten Programme die Umgebungsvariable 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 Generiert XLA-Programme, eines pro kompiliertem Cluster. Es ist äußerst hilfreich, diese beim Einreichen von XLA-Fehlerberichten anzuhängen.

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

  • module_XXXX.ptx Generierte PTX-Dateien.

Sie können die Grafik, die die Einbettung von XLA-Clustern visualisieren, innerhalb der TensorFlow-Grafik auch mit folgendem Befehl ausgeben:

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

Reproduzierbare Fehlerberichte

Ein Fehlerbericht ist viel einfacher zu reproduzieren, wenn er Dumps für die generierten XLA-Programme und die verwendete automatische Clustering-Einbettung enthält. Starten Sie folgenden Befehl, um sie für ein TensorFlow-Programm mit automatischem Clustering zu generieren:

$ 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"

Hängen Sie bei der Meldung von Fehlern den Inhalt des Verzeichnisses /tmp/generated an (siehe oben).

Versuchen Sie nach Möglichkeit, einen Fehler für ein einzelnes XLA-Programm zu isolieren. Verwenden Sie dazu run_hlo_module und führen Sie ihn iterativ in generierten Programmen aus.

Weitere Informationen

XLA-Frontends

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

  • JAX: Zusammensetzbare Transformationen von Python- und NumPy-Programmen
  • Julia: Die Julia-Sprache für wissenschaftliches Computing.
  • PyTorch: PyTorch-Framework
  • Nx: Numerische Rechenbibliothek für die Programmiersprache Elixir

Vorträge

XLA von TF mit jit_compile=True verwenden

XLA – Übersicht