Usamos OpenAI Triton para generar algunos de los kernels de GPU. Triton permite generar kernels de GPU rápidos para ciertas fusiones, pero debemos ajustar algunos parámetros para cada una de ellas.
Esto puede tardar mucho tiempo si hay muchas fusiones, por lo que proporcionamos una forma de cargar esos resultados de ajuste automático, a la vez que se ejecutan los otros pasos de compilación de forma normal. El ajuste automático de cachés aún es útil si hacemos algunos cambios: las fusiones que están presentes en la caché usarán la caché y las otras se ajustarán automáticamente de forma normal.
Recomendado: Directorio de caché
--xla_gpu_per_fusion_autotune_cache_dir=your/directory
Usa y mantén una caché de ajuste automático por fusión en el directorio determinado. Habrá un archivo por fusión distinta.
La principal ventaja de este enfoque es que puedes usar el mismo directorio de caché para varias ejecuciones de XLA (de diferentes modelos) y la caché crecerá con cada nueva fusión que se encuentre, lo que acelera las ejecuciones posteriores. También hay compatibilidad básica para ejecutar varias instancias de XLA con el mismo directorio de caché de forma simultánea.
XLA leerá los resultados existentes cuando sean necesarios y escribirá resultados nuevos después de que se determinen.
- El directorio debe existir antes de ejecutar XLA y debe ser de escritura.
- El usuario debe controlar la invalidación de la caché:
- Usa un directorio vacío si quieres comenzar con una caché vacía.
- El usuario debe realizar las verificaciones de versión de XLA:
- Si deseas usar cachés independientes para diferentes versiones de XLA, usa directorios diferentes.
La caché se desactiva de forma predeterminada (cuando no proporcionas el parámetro).
Limitación: No se garantiza que funcione bien en combinación con el otro método de almacenamiento en caché que se describe a continuación.
Alternativa: Carga o vuelca todos los resultados de un HLO determinado en un archivo
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, se volcará en formato protobuf binario.
En las pruebas
El ajuste automático persistente también se puede usar en 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 la que se ejecutan las pruebas.
Cómo hacer que una prueba use 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 del ajuste automático de la prueba, por ejemplo, especificando estos parámetros en el comando de prueba:
--test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto --test_sharding_strategy=disabledSe debe inhabilitar el fragmentación para obtener correctamente una sola caché de ajuste automático para todas las pruebas.
Luego, debemos subir esa caché a nuestro repositorio de código.
Luego, debemos agregar la caché a las dependencias de datos de nuestro objetivo 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 de 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 ralentizará. En este caso, deberíamos 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 XLA evoluciona y genera diferentes fusiones.