Usamos OpenAI Triton para generar algunos de los kernels de GPU. Triton permite generar kernels de GPU rápidos para ciertas fusiones, pero tenemos que ajustar algunos parámetros para cada una de ellas.
Esto puede demorar mucho tiempo si hay muchas fusiones, por lo que proporcionamos una manera de cargar los resultados de ajuste automático y, al mismo tiempo, ejecutar los otros pasos de compilación de forma normal. El ajuste automático de cachés sigue siendo útil si hacemos algunos cambios: las fusiones que están presentes en la caché usarán la caché y las otras se ajustarán de forma automática con normalidad.
Los resultados del ajuste automático se pueden volcar o cargar con estos parámetros:
--xla_gpu_dump_autotune_results_to=
--xla_gpu_load_autotune_results_from=
Si especificamos un archivo .txt o .textproto, la caché se volcará en formato textproto, de lo contrario, en formato protobuf binario.
En pruebas
El ajuste automático persistente también se puede usar en las pruebas. Se recomienda usarlo si las pruebas son muy grandes, en especial si el rendimiento del entorno de pruebas es limitado.
Solo funciona bien si la caché de ajuste automático contiene resultados generados en el mismo tipo de GPU en el que se ejecutan las pruebas.
Cómo hacer una prueba con el ajuste automático persistente
Por ahora, supongamos que la prueba en cuestión siempre usa el mismo tipo de GPU.
Debemos exportar los resultados de ajuste automático de la prueba, por ejemplo, especificando estos parámetros al comando de prueba:
--test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto --test_sharding_strategy=disabledLa fragmentación debe estar inhabilitada a fin de obtener correctamente una sola caché de ajuste automático para todas las pruebas.
Luego, tenemos que subir esa caché a nuestro repositorio de código.
Luego, tenemos que agregar la caché a las dependencias de datos del destino de prueba y cargarla con una variable de entorno.
data = ["test_autotune_cache.textproto"], env = {"XLA_FLAGS": "--xla_gpu_load_autotune_results_from=" + "$(execpath test_autotune_cache.textproto)"},(está bien usar la fragmentación en pruebas que carguen resultados del ajuste automático).
Consulta también las pruebas de ejemplo en 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
Obsolescencia de la caché
Si se realizan muchos cambios en un modelo, es posible que la caché ya no contenga todas las fusiones, por lo que la prueba se volverá más lenta. En este caso, tendríamos que volver a generar la caché de ajuste automático.
Si comenzamos a usar un nuevo tipo de GPU para ejecutar las pruebas, se aplica lo mismo.
La caché también puede volverse obsoleta si el compilador de XLA evoluciona y genera diferentes fusiones.