Zrzut obliczeń HLO

Zrzut HLO to tekstowa reprezentacja modułów HLO na różnych etapach obliczeń. Jest przydatny do debugowania i często trzeba go uwzględniać w raportach o błędach. Jest to zwykle czytelny dla człowieka plik tekstowy, który zawiera listę instrukcji HLO i ich właściwości. Czasami moduły HLO są zrzucane w postaci:

  • HloProto: pliki bufora protokołu, które są bardziej uporządkowanym formatem czytelnym dla komputera.
  • HloSnapshot: moduł HLO wraz z danymi wejściowymi. W przypadku ponownego odtwarzania HLO czasami potrzebne są rzeczywiste dane wejściowe przekazywane do danego obliczenia, a nie dane losowe.

Możesz używać flag XLA, aby określać zrzuty i je uzyskiwać. W większości przypadków możesz ustawić go za pomocą zmiennej środowiskowej. JAX oferuje też programowy sposób drukowania zrzutu HLO.

Lokalne wykonanie

Używanie zmiennych środowiskowych

Możesz ustawić zmienną środowiskową XLA_FLAGS z odpowiednimi flagami, aby uzyskać zrzuty. Działa to w przypadku JAX, TensorFlow i PyTorch/XLA.

Aby zrzucić moduły HLO i inne informacje do debugowania do określonego katalogu, uruchom program z flagą --xla_dump_to:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

Ścieżkami mogą być na przykład /tmp lub /tmp/xladump.

Domyślnie moduły HLO są zrzucane jako tekst na samym początku i na końcu potoku optymalizacji.

Możesz też wyraźnie określić format:

  1. zrzuty tekstu;
XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH"
  1. Protokoły HLO
XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH"
  1. Zrzuty HLO
XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH"
  1. Renderowanie grafu za pomocą serwera graphviz (działa dobrze tylko w przypadku małych grafów)
XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH"
  1. Renderowanie wykresu do pliku HTML (działa dobrze tylko w przypadku małych wykresów)
XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

W przypadku większych wykresów możesz użyć ikony interactive_graphviz, aby wizualizować części wykresu.

Zrzucanie konkretnych przebiegów pośrednich

Oprócz standardowych wstępnie zoptymalizowanych i ostatecznie zoptymalizowanych HLO możesz też zrzucać stan HLO po określonym etapie kompilacji.

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

Moduły HLO zostaną zrzucone dla przebiegów, których nazwy pasują do wyrażenia regularnego. Możesz na przykład obserwować HLO wynikające z przekazywania informacji związanych z podziałem SPMD za pomocą:

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

Aby zrzucać wynik po każdym przebiegu XLA (spowoduje to utworzenie wielu plików), możesz ustawić:

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

Opcje specyficzne dla JAX

Automatycznie w JAX

Zamiast przekazywać flagi lub zmienne środowiskowe, możesz też programowo zrzucać HLO za pomocą interfejsów API lower i compile JAX.

Lokalnie pobierz nieoptymalizowany, oryginalny, obniżony poziom HLO za pomocą tego polecenia:

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

Aby zrzucać dane do plików podczas etapów kompilacji HLO, określ:

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)

Zrzut jaxpr

jaxprs to pośrednia reprezentacja śladów programów w JAX. Aby to zrobić, ustaw zmienne środowiskowe:

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

Więcej informacji znajdziesz w dokumentacji JAX w sekcji Exporting and serializing staged-out computations: Debugging (Eksportowanie i serializowanie obliczeń etapowych: debugowanie).

Google Colab

Zmienne środowiskowe

W pierwszej wykonanej komórce notatnika (ponieważ zmienne środowiskowe i flagi wiersza poleceń są zwykle przetwarzane tylko raz, np.w momencie importowania modułu lub inicjowania backendu XLA) dodaj kod XLA_FLAGS opisany powyżej z os.environ, np.:

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

Spowoduje to przekazanie obliczeń do DIRECTORY_PATH, np. /tmp. W Colab otwórz przeglądarkę „Pliki” na pasku bocznym po lewej stronie, aby wyświetlić ten katalog i uzyskać do niego dostęp.

Możesz używać wszystkich flag wymienionych w sekcji Wykonywanie lokalne.

Opcje specyficzne dla JAX

Podobnie jak w przypadku wykonywania lokalnego, w przypadku interaktywnego sprawdzania na żywo możesz bezpośrednio wydrukować wstępnie zoptymalizowany kod HLO obliczeń:

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())

Możesz też bezpośrednio wydrukować zoptymalizowany HLO obliczeń:

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)

Zrzucanie wszystkich/małych obliczeń

Jeśli chcesz zobaczyć wszystko w zrzucie, w tym wszystkie małe kompilacje, ustaw zmienną środowiskową JAX:

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

Mozaika

Mosaic to kompilator dla backendu Pallas TPU i eksperymentalnego backendu Pallas GPU. Aby zrzucić obliczenia mozaiki, ustaw ten flagę:

--xla_mosaic_dump_to=/tmp/mosaic_dumps

Możesz też ustawić argumenty inicjowania TPU jako zmienną środowiskową:

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

Więcej informacji znajdziesz w dokumentacji JAX na temat Pallas i Mosaic.

Więcej z HLO Dumps

Znajdowanie odpowiednich obliczeń

Zazwyczaj wiele obliczeń jest odrzucanych. Zrzucane pliki mają nazwy zawierające „nazwę obliczeń” JAX, TensorFlow lub PyTorch/XLA, która jest podana w dziennikach, co ułatwia identyfikację odpowiednich plików HLO. Na przykład:

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

W przeciwnym razie możesz użyć ripgrep, aby szybko sprawdzić, w którym module znajdują się określone symbole lub obliczenia.

Wskazówka: w zgłoszeniach błędów uwzględnij 3 pliki zrzutu przed/po/przypisaniu bufora.

Konwersja HLO

Narzędzie hlo-opt, które może tłumaczyć formaty HLOProto i tekstowe. Jest to przydatne w sytuacjach, gdy masz jeden format, ale potrzebujesz drugiego do debugowania.

Dowiedz się, jak go używać: dokumentacja narzędzi XLA: hlo-opt.

Odtwórz ponownie

Możesz uruchomić (odtworzyć) zrzut obliczeń na określonym backendzie XLA z fałszywymi danymi lub zrzutami danych wejściowych. To wygodny sposób na odtwarzanie, iterowanie i debugowanie problemów w XLA.

Poniższe polecenia używają fałszywych danych. Jeśli masz zapisane migawki HLO, możesz je przekazać, a zostaną użyte dane z migawki. Aby nadal używać fałszywych danych podczas tworzenia migawki, przekaż flagę --force_fake_data.

Backend procesora:

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

Backend GPU:

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