Fortlaufendes Autotuning (nur GPU)

Wir verwenden OpenAI Triton, um einige der GPU-Kernel zu generieren. Triton ermöglicht das Generieren schneller GPU-Kernels für bestimmte Fusionen, aber wir müssen einige Parameter für jede solche Fusion anpassen.

Bei vielen Fusionen kann das sehr lange dauern. Daher bieten wir eine Möglichkeit, diese Ergebnisse der automatischen Optimierung zu laden, während die anderen Schritte der Kompilierung normal ausgeführt werden. Autotuning-Caches sind weiterhin nützlich, wenn wir einige Änderungen vornehmen: Die im Cache vorhandenen Fusionen verwenden den Cache und die anderen werden wie gewohnt automatisch optimiert.

--xla_gpu_per_fusion_autotune_cache_dir=your/directory

Verwende und verwalte einen Autotune-Cache pro Fusion im angegebenen Verzeichnis. Es gibt eine Datei pro Fusion.

Der Hauptvorteil dieses Ansatzes besteht darin, dass Sie dasselbe Cache-Verzeichnis für mehrere XLA-Ausführungen (verschiedener Modelle) verwenden können. Der Cache wird mit jeder neuen Fusion erweitert, was nachfolgende Ausführungen beschleunigt. Es gibt auch grundlegende Unterstützung für die gleichzeitige Ausführung mehrerer XLA-Instanzen mit demselben Cache-Verzeichnis.

XLA liest vorhandene Ergebnisse bei Bedarf und schreibt neue Ergebnisse, nachdem sie ermittelt wurden.

  • Das Verzeichnis muss vor dem Ausführen von XLA vorhanden sein und es muss beschreibbar sein.
  • Die Cache-Entwertung muss vom Nutzer durchgeführt werden:
    • Verwenden Sie ein leeres Verzeichnis, wenn Sie mit einem leeren Cache beginnen möchten.
  • XLA-Versionsprüfungen müssen vom Nutzer durchgeführt werden:
    • Wenn Sie separate Caches für verschiedene XLA-Versionen verwenden möchten, verwenden Sie bitte unterschiedliche Verzeichnisse.

Der Cache ist standardmäßig deaktiviert (wenn Sie den Parameter nicht angeben).

Einschränkung: Dies funktioniert in Kombination mit der anderen unten beschriebenen Caching-Methode möglicherweise nicht ordnungsgemäß.

Alternative: Laden oder Speichern aller Ergebnisse aus einem bestimmten HLO in eine Datei

Die Ergebnisse der automatischen Optimierung können mit den folgenden Parametern gedumpt/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 gedumpt, andernfalls im binären Protobuf-Format.

In Tests

Die dauerhafte automatische Abstimmung kann auch in Tests verwendet werden. Sie wird empfohlen, wenn die Tests sehr umfangreich sind, insbesondere wenn die Leistung der Testumgebung begrenzt ist.

Die Funktion funktioniert nur dann gut, wenn der Autotune-Cache Ergebnisse enthält, die auf derselben GPU generiert wurden, auf der die Tests ausgeführt werden.

Für einen Test die gespeicherte automatische Optimierung verwenden

Angenommen, für den betreffenden Test wird immer derselbe GPU-Typ verwendet.

  1. Wir müssen die Ergebnisse der automatischen Optimierung aus dem Test exportieren, z. B. indem wir diese Parameter für den Testbefehl angeben:

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

    Die Fragmentierung muss deaktiviert sein, damit für alle Tests ein einziger Autotune-Cache gespeichert wird.

  2. Anschließend müssen wir diesen Cache in unser Code-Repository hochladen.

  3. Dann müssen wir den Cache den Datenabhängigkeiten unseres Testziels hinzufügen und ihn mit einer Umgebungsvariablen laden.

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

    (Es ist in Ordnung, Sharding in Tests zu verwenden, bei denen Autotune-Ergebnisse geladen werden.)

Sehen Sie sich auch die Beispieltests unter 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

Wenn an einem Modell viele Änderungen vorgenommen werden, enthält der Cache möglicherweise nicht mehr alle Fusionen. In diesem Fall wird der Test langsamer. In diesem Fall müssten wir den Cache für die automatische Abstimmung neu generieren.

Dasselbe gilt, wenn wir für die Ausführung der Tests einen neuen GPU-Typ verwenden.

Der Cache kann auch veraltet werden, wenn der XLA-Compiler weiterentwickelt wird und andere Fusionen generiert.