Ajuste automático persistente (somente GPU)

Usamos o OpenAI Triton para gerar alguns kernels de GPU. O Triton permite gerar kernels de GPU rápidos para determinadas fusões, mas precisamos ajustar alguns parâmetros para cada uma delas.

Isso pode levar muito tempo se houver muitas fusões. Por isso, oferecemos uma maneira de carregar esses resultados de ajuste automático, enquanto executamos as outras etapas de compilação normalmente. Os caches de ajuste automático ainda são úteis quando fazemos algumas alterações: as fusões presentes no cache usarão o cache, e as outras serão ajustadas automaticamente normalmente.

--xla_gpu_per_fusion_autotune_cache_dir=your/directory

Use e mantenha um cache de autotune por fusão no diretório especificado. Haverá um arquivo para cada fusão distinta.

A principal vantagem dessa abordagem é que você pode usar o mesmo diretório de cache para várias execuções de XLA (de modelos diferentes), e seu cache vai crescer com cada nova fusão encontrada, acelerando as execuções subsequentes. Há também suporte básico para executar várias instâncias de XLA com o mesmo diretório de cache simultaneamente.

A XLA vai ler os resultados atuais quando eles forem necessários e gravar novos resultados depois que eles forem determinados.

  • O diretório precisa existir antes de executar o XLA e ser gravável.
  • A invalidação do cache precisa ser processada pelo usuário:
    • Use um diretório vazio se quiser começar com um cache vazio.
  • As verificações de versão do XLA precisam ser feitas pelo usuário:
    • Se quiser usar caches separados para versões diferentes do XLA, use diretórios diferentes.

O cache fica desativado por padrão (quando você não fornece o parâmetro).

Limitação: não há garantia de que isso funcione bem em combinação com o outro método de armazenamento em cache descrito abaixo.

Alternativa: carregar ou despejar todos os resultados de um determinado HLO em um arquivo

Os resultados da autoajuste podem ser descartados/carregados usando estes parâmetros:

--xla_gpu_dump_autotune_results_to=
--xla_gpu_load_autotune_results_from=

Se especificarmos um arquivo .txt ou .textproto, o cache será armazenado em formato textproto. Caso contrário, será no formato protobuf binário.

Em testes

A autoajuste persistente também pode ser usada em testes. Recomendamos usá-lo se os testes forem muito grandes, especialmente se o desempenho do ambiente de teste for limitado.

Ele só funciona bem se o cache de ajuste automático tiver resultados gerados no mesmo tipo de GPU em que os testes estão sendo executados.

Fazer um teste usar a autoajuste persistente

Por enquanto, vamos supor que o teste em questão sempre usa o mesmo tipo de GPU.

  1. Precisamos exportar os resultados do ajuste automático do teste, por exemplo, especificando estes parâmetros ao comando de teste:

    --test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto
--test_sharding_strategy=disabled

    A fragmentação precisa ser desativada para receber corretamente um único cache de ajuste automático para todos os testes.

  2. Em seguida, precisamos fazer o upload desse cache para o repositório de código.

  3. Em seguida, precisamos adicionar o cache às dependências de dados do nosso destino de teste e carregá-lo usando uma variável de ambiente.

    data = ["test_autotune_cache.textproto"],
env = {"XLA_FLAGS": "--xla_gpu_load_autotune_results_from=" +
                    "$(execpath test_autotune_cache.textproto)"},

    Não há problema em usar fragmentação em testes que carregam resultados de ajuste automático.

Consulte também os exemplos de testes em xla/service/gpu/tests/BUILD:

  • load_autotune_results_using_execpath_test
  • load_autotune_results_from_test_workspace_test
  • dump_autotune_results_to_test_outputs_test

Obsolescência do cache

Se muitas mudanças forem feitas em um modelo, é possível que o cache não contenha mais todas as fusões, e o teste ficará mais lento. Nesse caso, teríamos que gerar novamente o cache de ajuste automático.

Se começarmos a usar um novo tipo de GPU para executar os testes, o mesmo se aplica.

O cache também pode se tornar obsoleto se o compilador XLA evoluir e gerar diferentes fusões.