Réglage automatique persistant (GPU uniquement)

Nous utilisons OpenAI Triton pour générer certains des noyaux de GPU. Triton permet de générer des noyaux GPU rapides pour certaines fusions, mais nous devons ajuster certains paramètres pour chaque fusion de ce type.

Cela peut prendre beaucoup de temps s'il y a de nombreuses fusions. Nous fournissons donc un moyen de charger ces résultats d'auto-tuning, tout en exécutant les autres étapes de compilation normalement. Les caches d'ajustement automatique restent utiles si nous apportons quelques modifications: les fusions présentes dans le cache l'utiliseront, et les autres seront ajustées automatiquement.

--xla_gpu_per_fusion_autotune_cache_dir=your/directory

Utilisez et gérez un cache d'autotune par fusion dans le répertoire donné. Il y aura un fichier par fusion distincte.

L'avantage principal de cette approche est que vous pouvez utiliser le même répertoire de cache pour plusieurs exécutions XLA (de différents modèles). Votre cache augmentera à chaque nouvelle fusion rencontrée, ce qui accélérera les exécutions ultérieures. Il est également possible d'exécuter simultanément plusieurs instances XLA avec le même répertoire de cache.

XLA lit les résultats existants lorsqu'ils sont nécessaires et écrit de nouveaux résultats une fois qu'ils sont déterminés.

  • Le répertoire doit exister avant d'exécuter XLA et être accessible en écriture.
  • L'invalidation de cache doit être gérée par l'utilisateur :
    • Veuillez utiliser un répertoire vide si vous souhaitez commencer avec un cache vide.
  • Les vérifications de version XLA doivent être effectuées par l'utilisateur :
    • Si vous souhaitez utiliser des caches distincts pour différentes versions de XLA, veuillez utiliser des répertoires différents.

Le cache est désactivé par défaut (lorsque vous ne fournissez pas le paramètre).

Limite: il n'est pas garanti que cette méthode fonctionne bien avec l'autre méthode de mise en cache décrite ci-dessous.

Autre option: charger ou vider tous les résultats d'un HLO donné dans un seul fichier

Les résultats de l'autotuning peuvent être vidés/chargés à l'aide des paramètres suivants:

--xla_gpu_dump_autotune_results_to=
--xla_gpu_load_autotune_results_from=

Si nous spécifions un fichier .txt ou .textproto, le cache sera vidé au format textproto, sinon au format protobuf binaire.

En cours de test

L'autotuning persistant peut également être utilisé dans les tests. Il est recommandé de l'utiliser si les tests sont très volumineux, en particulier si les performances de l'environnement de test sont limitées.

Il ne fonctionne correctement que si le cache d'ajustement automatique contient des résultats générés sur le même type de GPU que celui sur lequel les tests sont exécutés.

Faire en sorte qu'un test utilise l'auto-réglage persistant

Pour l'instant, supposons que le test en question utilise toujours le même type de GPU.

  1. Nous devons exporter les résultats de l'ajustement automatique à partir du test, par exemple en spécifiant ces paramètres à la commande de test:

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

    Le fractionnement doit être désactivé pour obtenir un seul cache d'ajustement automatique pour tous les tests.

  2. Nous devons ensuite importer ce cache dans notre dépôt de code.

  3. Nous devons ensuite ajouter le cache aux dépendances de données de notre cible de test, puis le charger à l'aide d'une variable d'environnement.

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

    (Vous pouvez utiliser le sharding dans les tests qui chargent les résultats d'ajustement automatique.)

Veuillez également consulter les exemples de tests dans 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

Obsolescence du cache

Si de nombreuses modifications sont apportées à un modèle, il est possible que le cache ne contienne plus toutes les fusions, ce qui ralentit le test. Dans ce cas, nous devrions régénérer le cache d'ajustement automatique.

Si nous commençons à utiliser un nouveau type de GPU pour exécuter les tests, il en va de même.

Le cache peut également devenir obsolète si le compilateur XLA évolue et génère des fusions différentes.