Ce guide propose une sélection de flags XLA clés pour aider les utilisateurs à naviguer et à utiliser efficacement les fonctionnalités de XLA. Les sections suivantes détaillent les indicateurs qui peuvent avoir un impact significatif sur les performances d'exécution et l'utilisation de la mémoire. Si des problèmes, tels que des plantages, surviennent après l'activation d'un indicateur, il est recommandé de rétablir le paramètre par défaut et de créer un problème GitHub.
Indicateurs de performances
Les indicateurs suivants sont essentiels pour améliorer les performances d'exécution. Tester ces paramètres peut vous permettre d'améliorer considérablement vos performances.
| Option | Description | Valeurs par défaut | Valeurs suggérées | Valeurs des candidats |
|---|---|---|---|---|
| Pipeline 1. xla_should_allow_loop_variant_parameter_in_chain 2. xla_should_add_loop_invariant_op_in_chain 3. xla_tpu_enable_ici_ag_pipelining |
Ces trois indicateurs doivent être utilisés conjointement pour activer le pipeline collectif des opérations all-gather ICI(Interchip-Interconnect), ce qui crée davantage de possibilités d'exécution simultanée. | 1. xla_should_allow_loop_variant_parameter_in_chain=kDisabled 2. xla_should_add_loop_invariant_op_in_chain=kDisabled 3. xla_tpu_enable_ici_ag_pipelining=false |
1. xla_should_allow_loop_variant_parameter_in_chain=kEnabled 2. xla_should_add_loop_invariant_op_in_chain=kEnabled 3. xla_tpu_enable_ici_ag_pipelining=true |
1. xla_should_allow_loop_variant_parameter_in_chain=kDisabled/kEnabled/kAuto 2. xla_should_add_loop_invariant_op_in_chain=kDisabled/kEnabled/kAuto 3. xla_tpu_enable_ici_ag_pipelining=true/false |
v5e/Async xla_enable_async_all_gather xla_tpu_enable_async_collective_fusion xla_tpu_enable_async_collective_fusion_fuse_all_gather |
Ces trois indicateurs doivent être utilisés conjointement pour activer les opérations all-gather asynchrones sur v5e. | xla_enable_async_all_gather=kAuto xla_tpu_enable_async_collective_fusion=true xla_tpu_enable_async_collective_fusion_fuse_all_gather=true |
xla_enable_async_all_gather=kAuto xla_tpu_enable_async_collective_fusion=true xla_tpu_enable_async_collective_fusion_fuse_all_gather=true |
xla_enable_async_all_gather=kDisabled/kEnabled/kAuto xla_tpu_enable_async_collective_fusion=true/false xla_tpu_enable_async_collective_fusion_fuse_all_gather=true/false |
v5e/Async xla_tpu_enable_async_collective_fusion xla_tpu_enable_async_collective_fusion_fuse_all_reduce |
Ces deux options doivent être utilisées conjointement pour activer les opérations all-reduce asynchrones sur v5e. | xla_tpu_enable_async_collective_fusion=true xla_tpu_enable_async_collective_fusion_fuse_all_reduce=false |
xla_tpu_enable_async_collective_fusion=true xla_tpu_enable_async_collective_fusion_fuse_all_reduce=true |
xla_tpu_enable_async_collective_fusion=true/false xla_tpu_enable_async_collective_fusion_fuse_all_reduce=true/false |
Async xla_tpu_enable_async_all_to_all |
Cet indicateur active la communication asynchrone tout-à-tout. | xla_tpu_enable_async_all_to_all=false |
xla_tpu_enable_async_all_to_all=true |
xla_tpu_enable_async_all_to_all=true/false |
Lié à la latence xla_all_gather_latency_bound_threshold_in_bytes |
Cette option est destinée aux opérations All-Gather à latence limitée (c'est-à-dire de petite taille). L'activation de cette option déclenche des optimisations spécifiques qui peuvent réduire le temps d'exécution des all-gathers liés à la latence. Il est généralement utilisé dans les charges de travail d'inférence. | xla_all_gather_latency_bound_threshold_in_bytes=-1 (qui n'est pas activé) |
4~16Mb(i.e. 4~16 * 1024 * 1024) |
[0, 9223372036854775807] |
Lié à la latence xla_all_reduce_latency_bound_threshold_in_bytes |
Cette option est destinée aux opérations All-Gather à latence limitée (c'est-à-dire de petite taille). L'activation de cette option déclenche des optimisations spécifiques qui peuvent réduire le temps d'exécution pour les all-reduces liés à la latence. Il est généralement utilisé dans les charges de travail d'inférence. | xla_all_reduce_latency_bound_threshold_in_bytes=-1 (qui n'est pas activé) |
4~16Mb(i.e. 4~16 * 1024 * 1024) |
[0, 9223372036854775807] |
Lié à la latence xla_collective_permute_latency_bound_threshold_in_bytes |
Cette option est destinée aux opérations All-Gather à latence limitée (c'est-à-dire de petite taille). L'activation de cette option déclenche des optimisations spécifiques qui peuvent réduire le temps d'exécution des permutations collectives liées à la latence. Il est généralement utilisé dans les charges de travail d'inférence. | xla_collective_permute_latency_bound_threshold_in_bytes=-1 (qui n'est pas activé) |
4~16Mb(i.e. 4~16 * 1024 * 1024) |
[0, 9223372036854775807] |
Lié à la latence xla_all_to_all_latency_bound_threshold_in_bytes |
Cette option est destinée aux opérations All-Gather à latence limitée (c'est-à-dire de petite taille). L'activation de cette option déclenche des optimisations spécifiques qui peuvent réduire le temps d'exécution pour les opérations all-to-all liées à la latence. Il est généralement utilisé dans les charges de travail d'inférence. | xla_all_to_all_latency_bound_threshold_in_bytes=-1 (qui n'est pas activé) |
4~16Mb(i.e. 4~16 * 1024 * 1024) |
[0, 9223372036854775807] |
xla_enable_async_collective_permute |
Réécrit toutes les opérations de permutation collective dans leurs variantes asynchrones. Lorsque la valeur est définie sur auto, XLA peut activer automatiquement le collectif asynchrone en fonction d'autres configurations ou conditions. |
xla_enable_async_collective_permute=kAuto |
xla_enable_async_collective_permute=kAuto |
xla_enable_async_collective_permute=kAuto/kEnabled/kDisabled |
Indicateurs de mémoire
Les indicateurs listés ci-dessous sont fournis pour résoudre les problèmes liés à HBM. Ces valeurs ne doivent être ajustées que si vous rencontrez des erreurs de type "mémoire insuffisante" HBM lors de la compilation du modèle. Dans tous les autres cas, nous vous recommandons de conserver les valeurs par défaut, car les modifier pourrait nuire à vos performances.
| Option | Description | Valeurs par défaut | Valeurs suggérées | Valeurs des candidats |
|---|---|---|---|---|
Scheduler xla_latency_hiding_scheduler_rerun |
Ce paramètre ajuste le comportement du planificateur de masquage de la latence. Pour ce faire, elle réduit progressivement la limite de mémoire allouée à la planification à chaque réexécution du processus. | xla_latency_hiding_scheduler_rerun=1 |
xla_latency_hiding_scheduler_rerun=5 |
0~10(it doesn’t make much sense beyond 10 reruns) |
Fusion xla_tpu_rwb_fusion |
Cet indicateur permet les fusions de type "reduce+broadcast" et peut réduire l'utilisation de la mémoire. | xla_tpu_rwb_fusion=true |
xla_tpu_rwb_fusion=false |
xla_tpu_rwb_fusion=true/false |
Scheduler xla_memory_scheduler |
Cet indicateur spécifie l'algorithme que le planificateur de mémoire utilisera pour minimiser la consommation de mémoire. L'utilisation d'un algorithme plus avancé peut permettre d'obtenir un programme moins gourmand en mémoire, mais au prix d'un temps de compilation plus long. | xla_memory_scheduler=kDefault |
xla_memory_scheduler=kBrkga |
xla_memory_scheduler=kDefault/kList/kDfs/kPostOrder/kBrkga |
Scheduler xla_tpu_enable_latency_hiding_scheduler |
Cet indicateur active le planificateur de masquage de la latence, qui nous permet d'effectuer des opérations collectives asynchrones au lieu d'opérations synchrones. La désactivation réduit l'utilisation de la mémoire, mais entraîne une perte des gains de performances de ces opérations asynchrones. | xla_tpu_enable_latency_hiding_scheduler=true |
xla_tpu_enable_latency_hiding_scheduler=false |
xla_tpu_enable_latency_hiding_scheduler=true/false |
SPMD xla_jf_spmd_threshold_for_windowed_einsum_mib |
Cet indicateur définit le seuil inférieur de la taille minimale du point pour déclencher matmul collectif. Si vous définissez une valeur plus élevée, vous économiserez de la mémoire, mais vous perdrez des opportunités d'effectuer des matmul collectifs. | xla_jf_spmd_threshold_for_windowed_einsum_mib=-1 |
10Mb~1Gb (i.e. 10*1024*1024 ~ 1024*1024*1024) |
[0, 9223372036854775807] |
Autres options couramment utilisées
| Option | Type | Remarques |
|---|---|---|
xla_dump_to |
Chaîne (chemin d'accès) | Dossier dans lequel les fichiers HLO pré-optimisés et d'autres artefacts seront placés (voir Outils XLA). |
Indicateurs TPU XLA
| Option | Type | Remarques |
|---|---|---|
xla_tpu_enable_data_parallel_all_reduce_opt |
Booléen (vrai/faux) | Optimisation visant à augmenter les opportunités de chevauchement pour les réductions "all-reduce" DCN (mise en réseau des centres de données) utilisées pour le sharding parallèle des données. |
xla_tpu_data_parallel_opt_different_sized_ops |
Booléen (vrai/faux) | Permet le pipelining des opérations parallèles de données sur plusieurs itérations, même si la taille de leurs sorties ne correspond pas à ce qui peut être enregistré sur place dans les variables empilées. Peut augmenter la pression sur la mémoire. |
xla_tpu_spmd_rng_bit_generator_unsafe |
Booléen (vrai/faux) | Indique si le HLO RngBitGenerator doit être exécuté de manière partitionnée. Cette méthode n'est pas sûre si des résultats déterministes sont attendus avec différents partitionnements sur différentes parties du calcul. |
xla_tpu_megacore_fusion_allow_ags |
Booléen (vrai/faux) | Permet de fusionner toutes les collectes avec des convolutions/all-reduces. |
xla_tpu_enable_ag_backward_pipelining |
Booléen (vrai/faux) | Les pipelines all-gather (actuellement les all-gather à méga-échelle) sont exécutés à l'envers dans les boucles d'analyse. |
Indicateurs XLA du GPU
| Option | Type | Remarques |
|---|---|---|
xla_gpu_enable_latency_hiding_scheduler |
Booléen (vrai/faux) | Cet indicateur permet aux planificateurs de masquage de la latence de chevaucher efficacement la communication asynchrone avec le calcul. La valeur par défaut est "False". |
xla_gpu_enable_triton_gemm |
Booléen (vrai/faux) | Utilisez la multiplication matricielle basée sur Triton. |
xla_gpu_graph_level |
Indicateur (de 0 à 3) | Ancienne option permettant de définir le niveau du graphique GPU. Utilisez xla_gpu_enable_command_buffer dans de nouveaux cas d'utilisation. 0 = désactivé ; 1 = capturer les fusions et les memcpys ; 2 = capturer les gemms ; 3 = capturer les convolutions. |
xla_gpu_all_reduce_combine_threshold_bytes |
Entier (octets) | Ces indicateurs permettent de déterminer quand combiner plusieurs petites opérations AllGather / ReduceScatter / AllReduce en une seule grande opération AllGather / ReduceScatter / AllReduce afin de réduire le temps passé à la communication entre les appareils. Par exemple, pour les seuils AllGather / ReduceScatter sur une charge de travail basée sur Transformer, envisagez de les régler suffisamment haut pour combiner au moins le poids AllGather / ReduceScatter d'une couche Transformer. Par défaut, la valeur combine_threshold_bytes est définie sur 256. |
xla_gpu_all_gather_combine_threshold_bytes |
Entier (octets) | Consultez xla_gpu_all_reduce_combine_threshold_bytes ci-dessus. |
xla_gpu_reduce_scatter_combine_threshold_bytes |
Entier (octets) | Consultez xla_gpu_all_reduce_combine_threshold_bytes ci-dessus. |
xla_gpu_enable_pipelined_all_gather |
Booléen (vrai/faux) | Activez le pipelining des instructions all-gather. |
xla_gpu_enable_pipelined_reduce_scatter |
Booléen (vrai/faux) | Active le pipelining des instructions reduce-scatter. |
xla_gpu_enable_pipelined_all_reduce |
Booléen (vrai/faux) | Active le pipelining des instructions all-reduce. |
xla_gpu_enable_while_loop_double_buffering |
Booléen (vrai/faux) | Activez la double mise en mémoire tampon pour la boucle while. |
xla_gpu_enable_all_gather_combine_by_dim |
Booléen (vrai/faux) | Combiner les opérations all-gather avec la même dimension de collecte ou quelle que soit leur dimension. |
xla_gpu_enable_reduce_scatter_combine_by_dim |
Booléen (vrai/faux) | Combinez les opérations reduce-scatter avec la même dimension ou indépendamment de leur dimension. |