Dump HLO Computations

Un dump HLO è una rappresentazione testuale dei moduli HLO in diverse fasi del calcolo. È utile per il debug e spesso devi includerlo nelle segnalazioni di bug. In genere si tratta di un file di testo leggibile che elenca le istruzioni HLO e le relative proprietà. A volte, i moduli HLO vengono scaricati come:

  • HloProto:file buffer di protocollo, che sono un formato più strutturato e leggibile dal computer.
  • HloSnapshot: modulo HLO più i relativi input. Per riprodurre gli HLO, a volte sono necessari gli input effettivi forniti a un determinato calcolo anziché dati casuali.

Puoi utilizzare i flag XLA per specificare e ottenere i dump. Nella maggior parte dei casi, puoi impostarlo con una variabile di ambiente. JAX offre anche un modo programmatico per stampare il dump HLO.

Esecuzione locale

Utilizzo delle variabili di ambiente

Puoi impostare la variabile di ambiente XLA_FLAGS con i flag necessari per ottenere i dump. Questa funzionalità è valida per JAX, TensorFlow e PyTorch/XLA.

Per eseguire il dump dei moduli HLO e di altre informazioni di debug in una directory specifica, esegui il programma con il flag --xla_dump_to:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

Ad esempio, puoi utilizzare /tmp o /tmp/xladump come percorsi.

Per impostazione predefinita, i moduli HLO vengono scaricati come testo all'inizio e alla fine della pipeline di ottimizzazione.

Puoi anche specificare esplicitamente il formato:

  1. Dump di testo
XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH"
  1. HLO protos
XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH"
  1. Snapshot HLO
XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH"
  1. Rendering del grafico con il server graphviz (funziona bene solo per grafici di piccole dimensioni)
XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH"
  1. Rendering del grafico in file HTML (funziona bene solo per grafici di piccole dimensioni)
XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

Per i grafici più grandi, puoi utilizzare interactive_graphviz per visualizzare parti del grafico.

Passaggi intermedi specifici per lo scarico

Oltre agli HLO pre-ottimizzati / ottimizzati finali standard, puoi anche scaricare lo stato degli HLO dopo un particolare passaggio del compilatore.

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

I moduli HLO verranno scaricati per le tessere i cui nomi corrispondono all'espressione regolare (regex). Ad esempio, puoi osservare gli HLO risultanti dai passaggi relativi al partizionamento SPMD con:

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

Per eseguire il dump del risultato dopo ogni passaggio XLA (in questo modo verranno generati molti file), puoi impostare:

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

Opzioni specifiche di JAX

In modo programmatico in JAX

Anziché passare flag o variabili di ambiente, puoi anche eseguire il dump di HLO in modo programmatico utilizzando le API lower e compile di JAX.

Recupera localmente l'HLO originale non ottimizzato con:

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

Per il dump nei file durante le passate di compilazione HLO, specifica:

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)

Dump jaxprs

Gli jaxpr sono la rappresentazione intermedia di JAX per le tracce dei programmi. Per eseguire il dump, imposta le variabili di ambiente:

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

Scopri di più nella documentazione di JAX su Esportazione e serializzazione di calcoli in staging: debug.

Google Colab

Variabili di ambiente

Nella prima cella eseguita del notebook (perché le variabili di ambiente e i flag della riga di comando vengono in genere elaborati una sola volta, ad esempio al momento dell'importazione del modulo o dell'inizializzazione del backend XLA), aggiungi XLA_FLAGS come descritto sopra con os.environ, ad esempio:

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

In questo modo, il calcolo verrà eseguito su DIRECTORY_PATH, ad esempio /tmp. In Colab, vai al browser "File" nella barra laterale a sinistra per visualizzare e accedere a questa directory.

Puoi utilizzare tutti i flag menzionati nella sezione Esecuzione locale.

Opzioni specifiche di JAX

Simile all'esecuzione locale; per l'introspezione interattiva in tempo reale puoi stampare direttamente l'HLO pre-ottimizzato di un calcolo:

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

Puoi anche stampare direttamente l'HLO ottimizzato di un calcolo:

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)

Dumping di tutti i calcoli/di piccole dimensioni

Se vuoi visualizzare tutti i contenuti di un dump, comprese tutte le piccole compilazioni, imposta la variabile di ambiente JAX:

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

Mosaico

Mosaic è un compilatore per il backend Pallas TPU e il backend Pallas GPU sperimentale. Per eseguire il dump del calcolo del mosaico, imposta il seguente flag:

--xla_mosaic_dump_to=/tmp/mosaic_dumps

In alternativa, imposta gli argomenti di inizializzazione della TPU come variabile di ambiente:

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

Per saperne di più, consulta la documentazione di JAX su Pallas e Mosaic.

Altro con HLO Dumps

Trovare il calcolo giusto

Di solito, molti calcoli vengono eliminati. I file di dump sono denominati in modo esplicito con il "nome del calcolo" JAX, TensorFlow o PyTorch/XLA indicato nei log, il che semplifica l'identificazione dei file HLO pertinenti. Ad esempio:

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

In caso contrario, puoi utilizzare ripgrep per identificare rapidamente il modulo che contiene simboli o calcoli particolari.

Suggerimento:includi nei report sui bug i tre file di dump prima/dopo/assegnazione buffer di interesse.

Conversione HLO

Uno strumento chiamato hlo-opt che può tradurre tra i formati HLOProto e di testo. È utile nei casi in cui hai un formato, ma hai bisogno dell'altro per il debug.

Scopri come utilizzarlo: Documentazione di XLA Tooling: hlo-opt.

Guarda di nuovo

Puoi eseguire (riprodurre) i calcoli di cui è stato eseguito il dump su un backend XLA specificato con dati fittizi o snapshot di input. Questo è un modo conveniente per riprodurre, iterare e eseguire il debug dei problemi in XLA.

I seguenti comandi utilizzano dati fittizi. Se hai salvato gli snapshot HLO, puoi passarli e verranno utilizzati i dati dello snapshot. Per continuare a utilizzare dati fittizi durante l'esecuzione dello snapshot, passa il flag --force_fake_data.

Backend CPU:

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