Réglage automatique persistant (GPU uniquement)

Nous utilisons OpenAI Triton pour générer certains noyaux GPU. Triton permet de générer des noyaux GPU rapides pour certaines fusions, mais nous devons régler certains paramètres pour chacune de ces fusions.

Cela peut prendre beaucoup de temps si les fusions sont nombreuses. Nous fournissons donc un moyen de charger ces résultats d'autoréglage, tout en continuant à exécuter les autres étapes de compilation normalement. Les caches de réglage automatique sont toujours utiles si nous apportons quelques modifications: les fusions présentes dans le cache utiliseront le cache, et les autres seront réglés automatiquement normalement.

Les résultats de l'autoréglage 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, ou au format protobuf binaire.

En cours de test

Vous pouvez également utiliser le réglage automatique persistant 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.

Cela ne fonctionne bien que si le cache de réglage automatique contient des résultats générés sur le même type de GPU que celui où les tests sont exécutés.

Effectuer un test à l'aide du réglage automatique 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 du réglage automatique du test, par exemple en spécifiant ces paramètres dans 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
    

    La segmentation doit être désactivée pour obtenir correctement un seul cache de réglage 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 la cible de test et 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 la segmentation dans les tests qui chargent les résultats de l'autoréglage.)

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 générer à nouveau 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 différentes fusions.