Usamos o OpenAI Triton para gerar alguns dos 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 dessas fusões.
Isso pode levar muito tempo se houver muitas fusões. Portanto, 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 serão úteis se fizermos algumas alterações: as fusões presentes no cache usarão o cache, e as outras serão ajustadas automaticamente normalmente.
Os resultados do ajuste automático podem ser despejados/carregados com 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á despejado no formato textproto, caso contrário, no formato protobuf binário.
Em testes
O ajuste automático persistente também pode ser usado em testes. É recomendável usá-lo se os testes forem muito grandes, especialmente se o desempenho do ambiente de teste for limitado.
Ele só funcionará bem se o cache de ajuste automático contiver resultados gerados no mesmo tipo de GPU em que os testes estão sendo executados.
Como fazer um teste usar o ajuste automático persistente
Por enquanto, vamos supor que o teste em questão sempre use o mesmo tipo de GPU.
É preciso exportar os resultados do ajuste automático do teste, por exemplo, especificando esses parâmetros para o comando de teste:
--test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto --test_sharding_strategy=disabledA fragmentação precisa ser desativada para receber corretamente um único cache de ajuste automático para todos os testes.
Em seguida, temos que fazer upload do cache para nosso repositório de código.
Em seguida, precisamos adicionar o cache às dependências de dados do 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 a fragmentação em testes que carregam resultados do ajuste automático.
Consulte também os testes de exemplo 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 alterações 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 poderá se tornar obsoleto se o compilador XLA evolui e gera diferentes fusões.