XLA: Compiler für maschinelles Lernen optimieren

OpenXLA ist ein domänenspezifischer Compiler für die lineare Algebra, mit dem TensorFlow-Modelle möglicherweise ohne Änderungen am Quellcode beschleunigt werden können.

Einführung

Wenn ein TensorFlow-Programm ausgeführt wird, werden alle Vorgänge einzeln vom TensorFlow-Ausführer ausgeführt. Für jeden TensorFlow-Vorgang gibt es eine vorkompilierte GPU-Kernelimplementierung, an die der Executor sendet.

XLA bietet einen alternativen Modus zum Ausführen von Modellen: Es kompiliert den TensorFlow-Graphen in eine Sequenz von Rechenkernen, die speziell für das jeweilige Modell generiert wurden. Da diese Kernel für das Modell spezifisch sind, können sie modellspezifisch optimiert werden. Sehen wir uns zum Beispiel eine Optimierung an, die XLA im Kontext einer einfachen TensorFlow-Berechnung vornimmt:

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

Ohne XLA werden drei Kernel gestartet: einer für die Multiplikation, einer für die Addition und einer für die Reduktion. XLA kann den Graphen jedoch so optimieren, dass das Ergebnis in einem einzigen Kernelstart berechnet wird. Dazu werden Addition, Multiplikation und Reduktion in einem einzigen GPU-Kernel „zusammengeführt“. Außerdem werden bei dieser kombinierten Operation die von y*z und x+y*z erzeugten Zwischenwerte nicht in den Arbeitsspeicher geschrieben. Stattdessen werden die Ergebnisse dieser Zwischenberechnungen direkt an die Nutzer gestreamt, während sie vollständig in GPU-Registern bleiben. Die Fusion ist die wichtigste Optimierungsmethode von XLA. Die Arbeitsspeicherbandbreite ist in der Regel die knappste Ressource bei Hardwarebeschleunigern. Daher ist das Entfernen von Arbeitsspeichervorgängen eine der besten Möglichkeiten, die Leistung zu verbessern.

XLA für TensorFlow-Modelle aktivieren

Explizite Kompilierung mit tf.function(jit_compile=True)

Die API zur expliziten Kompilierung bietet eine detaillierte Auswahl, welche Funktionen kompiliert werden sollen. Beispielsweise wird die folgende TensorFlow-Funktion, die das MNIST-Training durchfü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 die Semantik muss-kompiliert: Entweder wird die gesamte Funktion mit XLA kompiliert oder es wird eine errors.InvalidArgumentError-Ausnahme geworfen. XLA kann derzeit keine Funktionen kompilieren, bei denen die Dimensionen nicht abgeleitet werden können, d. h., wenn es nicht möglich ist, die Dimensionen aller Tensoren ohne Ausführen der gesamten Berechnung abzuleiten. Die folgende Funktion wird beispielsweise nicht kompiliert:

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

Die Formen können jedoch zwischen den Läufen 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]))

In der Anleitung finden Sie ein detaillierteres Beispiel und ein 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)

Verwendung mit einer verteilten 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 des automatischen Clusterings. Dabei werden automatisch Cluster (verbundene Subgraphen) innerhalb der TensorFlow-Funktionen gefunden, die mit XLA kompiliert und ausgeführt werden können. Das automatische Clustern auf der GPU kann durch Festlegen der Umgebungsvariablen TF_XLA_FLAGS aktiviert werden:

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

Das automatische Clustern ist derzeit für GPU-Arbeitslasten optimiert, kann aber auch für die CPU aktiviert werden, indem zusätzlich das Flag --tf_xla_cpu_global_jit verwendet wird:

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

Ein detailliertes Beispiel finden Sie in der Colab-Anleitung zum automatischen Clustern.

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

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

Kompilierte Programme prüfen

XLA bietet Funktionen zur Inspektion der generierten Programme. Verwenden Sie die Umgebungsvariable XLA_FLAGS, um die generierten Programme zu dumpen:

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

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

  • module_XXXX.*_optimizations.txt Generierte XLA-Programme, eines pro kompiliertem Cluster. Es ist sehr hilfreich, diese beim Einreichen von XLA-Fehlerberichten anzuhängen.

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

  • module_XXXX.ptx Generierte PTX-Dateien.

Sie können auch das Diagramm mit der Einbettung von XLA-Clustern in die TensorFlow-Grafik mit folgendem Befehl dumpen:

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

Reproduzierbare Fehlerberichte

Ein Fehlerbericht lässt sich viel einfacher reproduzieren, wenn er Dumps für die generierten XLA-Programme und die verwendete automatische Clustering-Embedding enthält. Wenn Sie sie für ein TensorFlow-Programm generieren möchten, das mit automatischem Clustering ausgeführt wird, führen Sie Folgendes aus:

$ 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 beim Melden von Fehlern den Inhalt des Verzeichnisses /tmp/generated (siehe oben) an.

Versuchen Sie nach Möglichkeit, einen Fehler auf ein einzelnes XLA-Programm zu begrenzen, indem Sie run_hlo_module verwenden und es iterativ auf generierte Programme ausführen.

Weitere Informationen

XLA-Frontends

Neben TensorFlow können XLA-Programme auch mit folgenden Tools generiert werden:

  • JAX: Kombinierbare Transformationen von Python- und NumPy-Programmen
  • Julia: Die Julia-Programmiersprache für wissenschaftliche Berechnungen
  • PyTorch: PyTorch-Framework
  • Nx: Bibliothek für numerische Berechnungen für die Programmiersprache Elixir

Vorträge

XLA von TF mit jit_compile=True verwenden

XLA – Übersicht