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 Quellcodeänderungen beschleunigen kann.

Die Ergebnisse sind Verbesserungen in der Geschwindigkeit und Speicherverbrauch: zB in BERT MLPerf Vorlage mit 8 Volta V100 GPUs XLA mit hat eine 7x Verbesserung Leistung erreicht ~ und ~ 5x Losgröße Verbesserung:

Einführung

Wenn ein TensorFlow-Programm ausgeführt wird, werden alle Operationen einzeln vom TensorFlow-Executor ausgeführt. Jeder TensorFlow-Vorgang verfügt über 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 gegebene Modell generiert wurden. Da diese Kernel für das Modell einzigartig sind, können sie modellspezifische Informationen zur Optimierung nutzen. Sehen 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)

Ohne XLA ausgeführt, startet der Graph drei Kernel: einen für die Multiplikation, einen für die Addition und einen für die Reduktion. XLA kann den Graphen jedoch so optimieren, dass das Ergebnis in einem einzigen Kernelstart berechnet wird. Dies geschieht durch "Verschmelzen" der Addition, Multiplikation und Reduktion in einem einzigen GPU-Kernel. Darüber hinaus ist diese verschmolzenen Schreiboperation nicht die Zwischenwerte durch hergestellt aus y*z und x+y*z in den Speicher; stattdessen "streamt" es die Ergebnisse dieser Zwischenberechnungen direkt an ihre Benutzer, während sie vollständig in GPU-Registern gehalten werden. Fusion ist die wichtigste Optimierung von XLA. Die Speicherbandbreite ist in der Regel die knappste Ressource bei Hardwarebeschleunigern, daher ist das Entfernen von Speichervorgängen eine der besten Möglichkeiten, die Leistung zu verbessern.

XLA für TensorFlow-Modelle aktivieren

Explizite Kompilation mit tf.function(jit_compile=True)

Die API für explizite Kompilierung bietet eine fein abgestufte Steuerung für die Auswahl, welche Funktionen kompiliert werden sollen. Zum Beispiel 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 must-Kompilierung Semantik: entweder die gesamte Funktion wird mit XLA kompiliert oder eine errors.InvalidArgumentError Ausnahme ausgelöst. XLA kann derzeit nicht Funktionen zusammenstellen, wo Dimensionen sind nicht inferrable: Das heißt, wenn es nicht möglich ist , die Abmessungen aller Tensoren zu schließen , ohne die gesamte Berechnung ausgeführt wird . 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]))

Siehe das Tutorial colab ein ausführlichere Verwendungsbeispiel, und ein Tutorial - jit_compile=True Video auf jit_compile=True Nutzung.

Auto-Clustering

Eine einfache Art und Weise ohne Änderungen verwendet XLA in TensorFlow Modellen starten Auto-Clustering zu ermöglichen, die Cluster (verbundene Subgraphen) innerhalb der TensorFlow Funktionen findet automatisch die kompiliert werden kann , und unter Verwendung von XLA ausgeführt. Auto-Clustering auf GPU kann durch Einstellung der aktiviert sein TF_XLA_FLAGS Umgebungsvariable:

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

Auto-Clustering ist derzeit für GPU - Workloads optimiert, aber es kann auch durch zusätzlich mit der Flagge auf CPU aktiviert sein --tf_xla_cpu_global_jit :

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

Für ein detailliertes Anwendungsbeispiel des Auto-Clustering Tutorial colab .

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

Sie können auch einen Standalone verwenden tfcompile - Tool, das TensorFlow Graph in ausführbaren Code konvertiert (für x86-64 nur CPU).

Kompilierte Programme prüfen

XLA bietet Introspection-Funktionen, mit denen Sie die generierten Programme überprüfen können. Um die erzeugten Programme Dump, verwenden Sie die Umgebungsvariable XLA_FLAGS :

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

Nachdem das Dumping durchgeführt wird, können Sie die folgenden Dateien finden in /tmp/generated :

  • module_XXXX.*_optimizations.txt generiert XLA - Programme , eine pro jedem Cluster zusammengestellt. Diese beim Senden von XLA-Fehlerberichten anzuhängen ist äußerst hilfreich!

  • module_XXXX.ir-*.ll Erzeugen Dateien in LLVM Zwischendarstellung, mit NVPTX intrinsics.

  • module_XXXX.ptx generiert PTX - Dateien.

Sie können das Diagramm auch ausgeben, um die Einbettung von XLA-Clustern in das TensorFlow-Diagramm zu visualisieren mit:

$ 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 Auto-Clustering-Einbettung enthält. Um sie für ein TensorFlow-Programm zu generieren, das mit Auto-Clustering ausgeführt wird, starten Sie:

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

Wenn Sie Fehler Einreichung, befestigen Sie die Inhalte des /tmp/generated Verzeichnis (oben zitiert).

Wenn möglich, versuchen , einen Fehler zu einem einzigen XLA Programm zu isolieren , indem die Verwendung von replay_computation und es läuft iterativ generierten Programme.

Weiterlesen

XLA-Frontends

Abgesehen von TensorFlow können XLA-Programme generiert werden durch:

  • JAX : Composable Transformationen von Python + NumPy Programme
  • Julia : Die Julia Sprache für wissenschaftliches Rechnen
  • PyTorch : PyTorch Rahmen
  • Nx : Numerische Berechnung Bibliothek für die Programmiersprache Elixir

Gespräche

Mit XLA von TF mit jit_compile=True

XLA-Übersicht