Indicazioni sui flag XLA

Questa guida offre una selezione curata di flag XLA chiave per aiutare gli utenti a navigare e utilizzare in modo efficace le funzionalità di XLA. Le seguenti sezioni descrivono in dettaglio i flag che possono influire in modo significativo sulle prestazioni di runtime e sull'utilizzo della memoria. Se si verificano problemi, ad esempio arresti anomali, dopo l'attivazione di un flag, è consigliabile ripristinare l'impostazione predefinita e creare un problema su GitHub.

Flag di rendimento

I seguenti flag sono fondamentali per migliorare le prestazioni di runtime. Sperimentare con queste impostazioni può portare a notevoli miglioramenti del rendimento.

Bandiera Descrizione Valori predefiniti Valori suggeriti Valori del candidato
Pipelining
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
Questi tre flag devono essere utilizzati insieme per abilitare il pipelining collettivo delle operazioni di raccolta ICI(Interchip-Interconnect), il che crea più opportunità di esecuzione sovrapposta. 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
Questi tre flag devono essere utilizzati insieme per attivare le operazioni di raccolta asincrona su 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
Questi due flag devono essere utilizzati insieme per attivare le operazioni di riduzione asincrona su 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
Questo flag abilita la comunicazione asincrona all-to-all. 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
Limitato dalla latenza
xla_all_gather_latency_bound_threshold_in_bytes
Questo flag è destinato alle operazioni all-gather vincolate alla latenza (ovvero di piccole dimensioni). L'attivazione di questa opzione attiva ottimizzazioni specifiche che possono ridurre il tempo di esecuzione per tutti i gather vincolati dalla latenza. In genere viene utilizzato nei carichi di lavoro di inferenza. xla_all_gather_latency_bound_threshold_in_bytes=-1
(non abilitato)
4~16Mb(i.e. 4~16 * 1024 * 1024) [0, 9223372036854775807]
Limitato dalla latenza
xla_all_reduce_latency_bound_threshold_in_bytes
Questo flag è destinato alle operazioni all-gather vincolate alla latenza (ovvero di piccole dimensioni). L'attivazione di questa opzione attiva ottimizzazioni specifiche che possono ridurre il tempo di esecuzione per tutte le riduzioni vincolate dalla latenza. In genere viene utilizzato nei carichi di lavoro di inferenza. xla_all_reduce_latency_bound_threshold_in_bytes=-1
(non abilitato)
4~16Mb(i.e. 4~16 * 1024 * 1024) [0, 9223372036854775807]
Limitato dalla latenza
xla_collective_permute_latency_bound_threshold_in_bytes
Questo flag è destinato alle operazioni all-gather vincolate alla latenza (ovvero di piccole dimensioni). L'attivazione di questa opzione attiva ottimizzazioni specifiche che possono ridurre il tempo di esecuzione per le permutazioni collettive vincolate dalla latenza. In genere viene utilizzato nei carichi di lavoro di inferenza. xla_collective_permute_latency_bound_threshold_in_bytes=-1
(non abilitato)
4~16Mb(i.e. 4~16 * 1024 * 1024) [0, 9223372036854775807]
Limitato dalla latenza
xla_all_to_all_latency_bound_threshold_in_bytes
Questo flag è destinato alle operazioni all-gather vincolate alla latenza (ovvero di piccole dimensioni). L'attivazione di questa opzione attiva ottimizzazioni specifiche che possono ridurre il tempo di esecuzione per l'operazione all-to-all vincolata dalla latenza. In genere viene utilizzato nei carichi di lavoro di inferenza. xla_all_to_all_latency_bound_threshold_in_bytes=-1
(non abilitato)
4~16Mb(i.e. 4~16 * 1024 * 1024) [0, 9223372036854775807]
xla_enable_async_collective_permute Riscrive tutte le operazioni collective-permute nelle loro varianti asincrone. Se impostato su auto, XLA può attivare automaticamente la raccolta asincrona in base ad altre configurazioni o condizioni. xla_enable_async_collective_permute=kAuto xla_enable_async_collective_permute=kAuto xla_enable_async_collective_permute=kAuto/kEnabled/kDisabled

Flag di memoria

I flag elencati di seguito vengono forniti per risolvere i problemi relativi all'HBM. Questi devono essere modificati solo se si verificano errori di "memoria insufficiente" HBM durante la compilazione del modello. In tutti gli altri scenari, sono consigliati i valori predefiniti, in quanto la loro modifica potrebbe influire negativamente sul rendimento.

Bandiera Descrizione Valori predefiniti Valori suggeriti Valori del candidato
Scheduler
xla_latency_hiding_scheduler_rerun
Questa impostazione regola il comportamento dello scheduler di occultamento della latenza. Funziona riducendo in modo incrementale il limite di memoria allocato per la pianificazione a ogni "nuova esecuzione" del processo. 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
Questo flag attiva le fusioni di tipo reduce+broadcast e può ridurre l'utilizzo di memoria. xla_tpu_rwb_fusion=true xla_tpu_rwb_fusion=false xla_tpu_rwb_fusion=true/false
Scheduler
xla_memory_scheduler
Questo flag specifica l'algoritmo che lo scheduler di memoria utilizzerà per ridurre al minimo il consumo di memoria. L'utilizzo di un algoritmo più avanzato potrebbe generare una pianificazione che consuma meno memoria, a costo di un tempo di compilazione più lungo. xla_memory_scheduler=kDefault xla_memory_scheduler=kBrkga xla_memory_scheduler=kDefault/kList/kDfs/kPostOrder/kBrkga
Scheduler
xla_tpu_enable_latency_hiding_scheduler
Questo flag attiva lo scheduler di occultamento della latenza, che consente di eseguire operazioni collettive asincrone anziché sincrone. La sua disattivazione riduce l'utilizzo della memoria a scapito dei miglioramenti delle prestazioni derivanti da queste operazioni asincrone. 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
Questo flag imposta la soglia inferiore della dimensione minima del punto per attivare matmul collettivo. Se lo imposti su un valore più alto, risparmierai memoria a scapito della perdita di opportunità di eseguire matmul collettivo. xla_jf_spmd_threshold_for_windowed_einsum_mib=-1 10Mb~1Gb (i.e. 10*1024*1024 ~ 1024*1024*1024) [0, 9223372036854775807]

Altri flag di uso comune

Bandiera Tipo Note
xla_dump_to Stringa (percorso file) La cartella in cui verranno inseriti i file HLO pre-ottimizzati e altri artefatti (vedi XLA Tools).

Flag XLA TPU

Bandiera Tipo Note
xla_tpu_enable_data_parallel_all_reduce_opt Booleano (true/false) Ottimizzazione per aumentare le opportunità di sovrapposizione per le riduzioni di DCN (data center networking) utilizzate per lo sharding parallelo dei dati.
xla_tpu_data_parallel_opt_different_sized_ops Booleano (true/false) Consente il pipelining di operazioni parallele sui dati in più iterazioni anche se le dimensioni dell'output non corrispondono a quelle che possono essere salvate in loco nelle variabili impilate. Può aumentare la pressione della memoria.
xla_tpu_spmd_rng_bit_generator_unsafe Booleano (true/false) Indica se eseguire l'HLO RngBitGenerator in modo partizionato, il che non è sicuro se si prevedono risultati deterministici con sharding diversi in diverse parti del calcolo.
xla_tpu_megacore_fusion_allow_ags Booleano (true/false) Consente di unire tutti i gather con le convoluzioni/tutte le riduzioni.
xla_tpu_enable_ag_backward_pipelining Booleano (true/false) Pipeline all-gather (attualmente all-gather su larga scala) all'indietro attraverso i cicli di scansione.

Flag XLA per GPU

Bandiera Tipo Note
xla_gpu_enable_latency_hiding_scheduler Booleano (true/false) Questo flag consente agli scheduler di nascondere la latenza per sovrapporre in modo efficiente la comunicazione asincrona al calcolo. Il valore predefinito è False.
xla_gpu_enable_triton_gemm Booleano (true/false) Utilizza la moltiplicazione di matrici basata su Triton.
xla_gpu_graph_level Flag (0-3) Il flag legacy per impostare il livello del grafico GPU. Utilizza xla_gpu_enable_command_buffer in nuovi casi d'uso. 0 = off; 1 = capture fusions and memcpys; 2 = capture gemms; 3 = capture convolutions.
xla_gpu_all_reduce_combine_threshold_bytes Numero intero (byte) Questi flag regolano il momento in cui combinare più operazioni AllGather / ReduceScatter / AllReduce di piccole dimensioni in un'unica operazione AllGather / ReduceScatter / AllReduce di grandi dimensioni per ridurre il tempo dedicato alla comunicazione tra dispositivi. Ad esempio, per le soglie AllGather / ReduceScatter in un carico di lavoro basato su Transformer, valuta la possibilità di impostarle su un valore sufficientemente alto da combinare almeno AllGather / ReduceScatter del peso di un livello Transformer. Per impostazione predefinita, combine_threshold_bytes è impostato su 256.
xla_gpu_all_gather_combine_threshold_bytes Numero intero (byte) Vedi xla_gpu_all_reduce_combine_threshold_bytes sopra.
xla_gpu_reduce_scatter_combine_threshold_bytes Numero intero (byte) Vedi xla_gpu_all_reduce_combine_threshold_bytes sopra.
xla_gpu_enable_pipelined_all_gather Booleano (true/false) Abilita il pipelining delle istruzioni all-gather.
xla_gpu_enable_pipelined_reduce_scatter Booleano (true/false) Attiva il pipelining delle istruzioni di riduzione della dispersione.
xla_gpu_enable_pipelined_all_reduce Booleano (true/false) Attiva il pipelining di tutte le istruzioni di riduzione.
xla_gpu_enable_while_loop_double_buffering Booleano (true/false) Abilita il doppio buffering per il ciclo while.
xla_gpu_enable_all_gather_combine_by_dim Booleano (true/false) Combina tutte le operazioni di raccolta con la stessa dimensione di raccolta o indipendentemente dalla dimensione.
xla_gpu_enable_reduce_scatter_combine_by_dim Booleano (true/false) Combina le operazioni di riduzione della dispersione con la stessa dimensione o indipendentemente dalla dimensione.