HLO-Berechnungen ausgeben

Ein HLO-Dump ist eine Textdarstellung der HLO-Module in verschiedenen Phasen der Berechnung. Sie ist nützlich für das Debugging und muss oft in Fehlerberichten enthalten sein. Dies ist in der Regel eine für Menschen lesbare Textdatei, in der die HLO-Anweisungen und ihre Eigenschaften aufgeführt sind. Manchmal werden HLO-Module so ausgegeben:

  • HloProto:Protocol Buffer-Dateien, die ein strukturierteres, maschinenlesbares Format sind.
  • HloSnapshot: HLO-Modul und seine Eingaben. Zum Wiedergeben von HLOs benötigen Sie manchmal die tatsächlichen Eingaben für eine bestimmte Berechnung anstelle von zufälligen Daten.

Sie können XLA-Flags verwenden, um Dumps anzugeben und abzurufen. In den meisten Fällen können Sie sie mit einer Umgebungsvariable festlegen. JAX bietet auch eine programmatische Möglichkeit, den HLO-Dump auszugeben.

Lokale Ausführung

Umgebungsvariablen verwenden

Sie können die Umgebungsvariable XLA_FLAGS mit den erforderlichen Flags festlegen, um Dumps zu erhalten. Das funktioniert für JAX, TensorFlow und PyTorch/XLA.

Wenn Sie HLO-Module und andere Debugging-Informationen in ein bestimmtes Verzeichnis ausgeben möchten, führen Sie Ihr Programm mit dem Flag --xla_dump_to aus:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

Sie können beispielsweise /tmp oder /tmp/xladump als Pfade verwenden.

Standardmäßig werden HLO-Module als Text ausgegeben, ganz am Anfang und am Ende der Optimierungspipeline.

Sie können das Format auch explizit angeben:

  1. Textauszüge
XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH"
  1. HLO-Protokolle
XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH"
  1. HLO-Snapshots
XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH"
  1. Grafikrendering mit Graphviz-Server (funktioniert nur gut für kleine Grafiken)
XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH"
  1. Grafik als HTML-Datei rendern (funktioniert nur gut für kleine Grafiken)
XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

Bei größeren Diagrammen können Sie interactive_graphviz verwenden, um Teile des Diagramms zu visualisieren.

Dump Specific Intermediate Passes

Zusätzlich zu den standardmäßigen voroptimierten / final optimierten HLOs können Sie auch den Status von HLOs nach einem bestimmten Compiler-Pass ausgeben.

XLA_FLAGS="--xla_dump_hlo_pass_re=regex --xla_dump_to=DIRECTORY_PATH"

HLO-Module werden für die Durchläufe ausgegeben, deren Namen mit dem regulären Ausdruck (Regex) übereinstimmen. Beispielsweise können Sie die HLOs, die sich aus Durchläufen im Zusammenhang mit der SPMD-Partitionierung ergeben, mit dem folgenden Befehl beobachten:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=spmd|propagation"

Wenn Sie das Ergebnis nach jedem XLA-Pass ausgeben möchten (dies führt zu vielen Dateien), können Sie Folgendes festlegen:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=.*"

JAX-spezifische Optionen

Programmatisch in JAX

Anstatt Flags oder Umgebungsvariablen zu übergeben, können Sie HLO auch programmatisch mit den JAX-APIs lower und compile ausgeben.

Rufen Sie das nicht optimierte Original-HLO lokal mit folgendem Befehl ab:

jax.jit(f).lower(*args).as_text('hlo')

Geben Sie Folgendes an, um während der HLO-Kompilierungsdurchläufe in Dateien zu schreiben:

compilation_args = {
    'xla_dump_to': DIRECTORY_PATH,
    'xla_dump_hlo_pass_re': 'spmd|propagation', # or some other pass filter
    ...
    }

jax.jit(f).lower(*args).compile(compilation_args)

jaxprs sichern

jaxprs sind die Zwischenrepräsentation von JAX für Programm-Traces. Um dies zu protokollieren, legen Sie die Umgebungsvariablen fest:

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

Weitere Informationen finden Sie in der JAX-Dokumentation unter Exporting and serializing staged-out computations: Debugging.

Google Colab

Umgebungsvariablen

Fügen Sie in die erste ausgeführte Zelle Ihres Notebooks (da Umgebungsvariablen und Befehlszeilenflags in der Regel nur einmal verarbeitet werden, z. B. beim Modulimport oder bei der Initialisierung des XLA-Back-Ends) die oben beschriebene XLA_FLAGS mit os.environ ein, z. B.:

import os
os.environ['XLA_FLAGS'] = "--xla_dump_to=DIRECTORY_PATH"

Dadurch wird die Berechnung in DIRECTORY_PATH ausgegeben, z. B. /tmp. Rufen Sie in Colab in der linken Seitenleiste den Dateibrowser auf, um dieses Verzeichnis aufzurufen und darauf zuzugreifen.

Sie können alle im Abschnitt „Lokale Ausführung“ genannten Flags verwenden.

JAX-spezifische Optionen

Ähnlich wie bei der lokalen Ausführung können Sie für die interaktive Live-Introspektion das voroptimierte HLO einer Berechnung direkt ausgeben:

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

c = jax.jit(f).lower(3.).compiler_ir('hlo')

print(c.as_hlo_text())

Sie können auch das optimierte HLO einer Berechnung direkt ausdrucken:

def optimized_HLO(f, *args, platform=None):
    print(jax.jit(f).lower(*args).compile().as_text())

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

optimized_HLO(f, 1.0)

Alle/kleine Berechnungen verwerfen

Wenn Sie alles in einem Dump sehen möchten, einschließlich aller kleinen Kompilierungen, legen Sie die JAX-Umgebungsvariable so fest:

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

Mosaik

Mosaic ist ein Compiler für das Pallas-TPU-Backend und das experimentelle Pallas-GPU-Backend. Legen Sie das folgende Flag fest, um die Berechnung des Mosaiks auszugeben:

--xla_mosaic_dump_to=/tmp/mosaic_dumps

Alternativ können Sie TPU-Initialisierungsargumente als Umgebungsvariable festlegen:

export LIBTPU_INIT_ARGS="--xla_mosaic_dump_to=/tmp/mosaic_dumps"

Weitere Informationen

Mehr mit HLO Dumps

Die richtige Berechnung finden

Normalerweise werden viele Berechnungen verworfen. Die exportierten Dateien werden explizit mit dem JAX-, TensorFlow- oder PyTorch/XLA-„Berechnungsnamen“ benannt, der in den Logs angegeben ist. So lassen sich die relevanten HLO-Dateien leicht identifizieren. Beispiel:

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

Andernfalls können Sie ripgrep verwenden, um schnell zu ermitteln, in welchem Modul sich bestimmte Symbole oder Berechnungen befinden.

Tipp:Fügen Sie Ihren Fehlerberichten die drei Dateien mit den Dumps vor, nach und während der Pufferzuweisung bei.

HLO-Conversion

Ein Tool namens hlo-opt, das zwischen HLOProto- und Textformaten übersetzen kann. Das ist nützlich, wenn Sie ein Format haben, aber das andere zum Debuggen benötigen.

Informationen zur Verwendung: XLA Tooling-Dokumentation: hlo-opt.

Nochmal

Sie können die gesicherten Berechnungen auf einem angegebenen XLA-Backend mit gefälschten Daten oder Eingabe-Snapshots ausführen (wiedergeben). So lassen sich Probleme in XLA einfach reproduzieren, iterieren und debuggen.

In den folgenden Befehlen werden gefälschte Daten verwendet. Wenn Sie HLO-Snapshots gespeichert haben, können Sie diese stattdessen übergeben. Dann werden die Daten aus dem Snapshot verwendet. Wenn Sie beim Ausführen des Snapshots weiterhin gefälschte Daten verwenden möchten, übergeben Sie das Flag --force_fake_data.

CPU-Backend:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=cpu
 /tmp/xladump/module_4561.before_optimizations.txt

GPU-Backend:

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=CUDA
 /tmp/xladump/module_4561.before_optimizations.txt