Wir verwenden OpenAI Triton, um einige der GPU-Kernel zu generieren. Mit Triton lassen sich schnelle GPU-Kernel für bestimmte Verschmelzungen generieren. Allerdings müssen für jede solche Verschmelzung einige Parameter angepasst werden.
Dies kann bei vielen Fusionen lange dauern. Daher bieten wir eine Möglichkeit, diese Autotuning-Ergebnisse zu laden, während die anderen Kompilierungsschritte normal ausgeführt werden. Die automatischen Abstimmungs-Caches sind trotzdem nützlich, wenn wir ein paar Änderungen vornehmen: Die im Cache vorhandenen Fusionen verwenden den Cache, die anderen werden normal automatisch abgestimmt.
Die Ergebnisse der AutoAbstimmung können mit folgenden Parametern in eine Dump-/Ladedatei geladen werden:
--xla_gpu_dump_autotune_results_to=
--xla_gpu_load_autotune_results_from=
Wenn wir eine TXT- oder Textproto-Datei angeben, wird der Cache im Textproto-Format ausgegeben, ansonsten im binären protobuf-Format.
In Tests
Die persistente automatische Einstellung kann auch in Tests verwendet werden. Es wird empfohlen, ihn zu verwenden, wenn die Tests sehr umfangreich sind, insbesondere wenn die Leistung der Testumgebung begrenzt ist.
Es funktioniert nur dann gut, wenn der Autotune-Cache Ergebnisse enthält, die mit demselben GPU-Typ generiert wurden, auf dem die Tests ausgeführt werden.
Bei Verwendung eines Tests wurde die automatische Einstellung beibehalten.
Nehmen wir zunächst einmal an, dass der fragliche Test immer denselben GPU-Typ verwendet.
Wir müssen die Autotune-Ergebnisse aus dem Test exportieren, indem wir beispielsweise die folgenden Parameter im Testbefehl angeben:
--test_env=XLA_FLAGS=--xla_gpu_dump_autotune_results_to=TEST_UNDECLARED_OUTPUTS_DIR/autotune_cache.textproto --test_sharding_strategy=disabledDie Fragmentierung muss deaktiviert sein, damit ein einzelner Autotune-Cache für alle Tests korrekt abgerufen wird.
Dann müssen wir diesen Cache in unser Code-Repository hochladen.
Anschließend müssen wir den Cache den Datenabhängigkeiten unseres Testziels hinzufügen und ihn mithilfe einer Umgebungsvariablen laden.
data = ["test_autotune_cache.textproto"], env = {"XLA_FLAGS": "--xla_gpu_load_autotune_results_from=" + "$(execpath test_autotune_cache.textproto)"},Die Fragmentierung kann in Tests verwendet werden, die Autotune-Ergebnisse laden.
Sehen Sie sich auch die Beispieltests in xla/service/gpu/tests/BUILD an:
- load_autotune_results_using_execpath_test
- load_autotune_results_from_test_workspace_test
- dump_autotune_results_to_test_outputs_test
Cache-Veralterung
Werden viele Änderungen an einem Modell vorgenommen, enthält der Cache möglicherweise nicht mehr alle Fusionen, sodass der Test langsamer wird. In diesem Fall müssten wir den Autotuning-Cache neu generieren.
Wenn wir zum Ausführen der Tests einen neuen GPU-Typ verwenden, gilt dasselbe.
Der Cache kann auch veraltet sein, wenn der XLA-Compiler weiterentwickelt wird und verschiedene Fusionen generiert.