Despejar cálculos de HLO

Um despejo de HLO é uma representação textual dos módulos de HLO em diferentes estágios da computação. Ele é útil para depuração e geralmente precisa ser incluído em relatórios de bugs. Normalmente, é um arquivo de texto legível por humanos que lista as instruções de HLO e as propriedades delas. Às vezes, os módulos HLO são despejados como:

  • HloProto:arquivos de buffer de protocolo, que são um formato mais estruturado e legível por máquina.
  • HloSnapshot: módulo HLO e entradas dele. Para reproduzir HLOs, às vezes é necessário usar as entradas reais transmitidas a um determinado cálculo em vez de dados aleatórios.

É possível usar flags XLA para especificar e receber dumps. Na maioria dos casos, é possível definir isso com uma variável de ambiente. O JAX também oferece uma maneira programática de imprimir o despejo de HLO.

Execução local

Como usar variáveis de ambiente

É possível definir a variável de ambiente XLA_FLAGS com as flags necessárias para receber dumps. Isso funciona para JAX, TensorFlow e PyTorch/XLA.

Para despejar módulos HLO e outras informações de depuração em um diretório específico, execute o programa com a flag --xla_dump_to:

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

Por exemplo, é possível usar /tmp ou /tmp/xladump como os caminhos.

Por padrão, isso despeja módulos HLO como texto, no início e no fim do pipeline de otimização.

Também é possível especificar o formato explicitamente:

  1. Despejos de texto
XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH"
  1. Protobufs HLO
XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH"
  1. Snapshots de HLO
XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH"
  1. Renderização de gráficos com o servidor graphviz (funciona bem apenas para gráficos pequenos)
XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH"
  1. Renderização de gráfico em arquivo HTML (funciona bem apenas para gráficos pequenos)
XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

Para gráficos maiores, use interactive_graphviz para visualizar partes do gráfico.

Despejar transmissões intermediárias específicas

Além dos HLOs padrão pré-otimizados / otimizados finais, também é possível despejar o estado dos HLOs após uma transmissão específica do compilador.

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

Os módulos HLO serão despejados para as transmissões cujos nomes correspondem à expressão regular (regex). Por exemplo, é possível observar os HLOs resultantes de transmissões relacionadas ao particionamento SPMD com:

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

Para despejar o resultado após cada transmissão XLA (isso vai gerar muitos arquivos), defina:

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

Opções específicas do JAX

De forma programática no JAX

Em vez de transmitir flags ou variáveis de ambiente, também é possível despejar programaticamente o HLO usando as APIs lower e compile do JAX.

Busque localmente o HLO reduzido original não otimizado com:

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

Para despejar em arquivos durante as transmissões de compilação de HLO, especifique:

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)

Despejar jaxprs

Os jaxprs são a representação intermediária do JAX para rastreamentos de programas. Para fazer isso, defina as variáveis de ambiente:

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

Saiba mais na documentação do JAX sobre Exportar e serializar cálculos em etapas: depuração.

Google Colab

Variáveis de ambiente

Na primeira célula executada do seu notebook (porque as variáveis de ambiente e as flags de linha de comando geralmente são processadas apenas uma vez, por exemplo, no momento da importação do módulo ou da inicialização do back-end XLA), adicione o XLA_FLAGS detalhado acima com os.environ. Por exemplo:

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

Isso vai despejar a computação em DIRECTORY_PATH, por exemplo, /tmp. No Colab, navegue até o navegador "Arquivos" na barra lateral esquerda para ver e acessar esse diretório.

Você pode usar todas as flags mencionadas na seção "Execução local".

Opções específicas do JAX

Semelhante à execução local. Para introspecção interativa e em tempo real, imprima diretamente o HLO pré-otimizado de uma computação:

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

Você também pode imprimir diretamente o HLO otimizado de um cálculo:

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)

Despejo de todos/pequenos cálculos

Se você quiser ver tudo em um despejo, incluindo todas as pequenas compilações, defina a variável de ambiente JAX:

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

Mosaico

O Mosaic é um compilador para o back-end de TPU do Pallas e o back-end experimental de GPU do Pallas. Para fazer um despejo do cálculo de mosaico, defina a seguinte flag:

--xla_mosaic_dump_to=/tmp/mosaic_dumps

Ou defina argumentos de inicialização da TPU como uma variável de ambiente:

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

Confira a documentação do JAX sobre Pallas e Mosaic para saber mais.

Mais com HLO Dumps

Como encontrar o cálculo certo

Normalmente, muitos cálculos são descartados. Os arquivos despejados são explicitamente nomeados com o "nome de computação" do JAX, TensorFlow ou PyTorch/XLA que são chamados nos registros, facilitando a identificação dos arquivos HLO relevantes. Exemplo:

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

Caso contrário, use ripgrep para identificar rapidamente qual módulo contém símbolos ou cálculos específicos.

Dica:inclua nos relatórios de bugs os três arquivos de despejo de antes/depois/atribuição de buffer de interesse.

Conversão de HLO

Uma ferramenta chamada hlo-opt que pode traduzir entre formatos HLOProto e de texto. Isso é útil quando você tem um formato, mas precisa do outro para depuração.

Saiba como usar: Documentação de ferramentas do XLA: hlo-opt.

Repetir

É possível executar (reproduzir) os cálculos despejados em um back-end XLA especificado com dados falsos ou snapshots de entrada. Essa é uma maneira conveniente de reproduzir, iterar e depurar problemas no XLA.

Os comandos a seguir usam dados falsos. Se você tiver salvo snapshots de HLO, poderá transmiti-los em vez disso, e os dados do snapshot serão usados. Para continuar usando dados falsos ao executar o snapshot, transmita a flag --force_fake_data.

Backend da CPU:

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

Backend da GPU:

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