Dieser Leitfaden bietet eine kuratierte Auswahl wichtiger XLA-Flags, um Nutzern zu helfen, die Funktionen von XLA effektiv zu nutzen. In den folgenden Abschnitten werden Flags beschrieben, die sich erheblich auf die Laufzeitleistung und die Speicherauslastung auswirken können. Sollten nach dem Aktivieren eines Flags Probleme wie Abstürze auftreten, empfiehlt es sich, zur Standardeinstellung zurückzukehren und ein GitHub-Problem zu erstellen.
Leistungsbezogene Meldungen
Die folgenden Flags sind entscheidend für die Verbesserung der Laufzeitleistung. Wenn Sie diese Einstellungen testen, kann sich die Leistung erheblich verbessern.
| Flag | Beschreibung | Standardwerte | Vorgeschlagene Werte | Kandidatenwerte |
|---|---|---|---|---|
| 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 |
Diese drei Flags sollten zusammen verwendet werden, um das kollektive Pipelining von ICI-All-Gather-Vorgängen(Interchip-Interconnect) zu ermöglichen. Dadurch ergeben sich mehr Möglichkeiten für die überlappende Ausführung. | 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 |
Diese drei Flags sollten zusammen verwendet werden, um asynchrone All-Gather-Vorgänge auf v5e zu aktivieren. | 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 |
Diese beiden Flags sollten zusammen verwendet werden, um asynchrone All-Reduce-Vorgänge auf v5e zu aktivieren. | 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 |
Dieses Flag aktiviert die asynchrone All-to-All-Kommunikation. | 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 |
Latenzgebunden xla_all_gather_latency_bound_threshold_in_bytes |
Dieses Flag ist für latenzgebundene (d.h. kleine) All-Gather-Vorgänge vorgesehen. Wenn Sie diese Option aktivieren, werden bestimmte Optimierungen ausgelöst, die die Ausführungszeit für latenzgebundene All-Gathers verkürzen können. Sie wird in der Regel für Inferenz-Arbeitslasten verwendet. | xla_all_gather_latency_bound_threshold_in_bytes=-1 (nicht aktiviert) |
4~16Mb(i.e. 4~16 * 1024 * 1024) |
[0, 9223372036854775807] |
Latenzgebunden xla_all_reduce_latency_bound_threshold_in_bytes |
Dieses Flag ist für latenzgebundene (d.h. kleine) All-Gather-Vorgänge vorgesehen. Wenn Sie diese Option aktivieren, werden bestimmte Optimierungen ausgelöst, die die Ausführungszeit für latenzgebundene All-Reduces reduzieren können. Sie wird in der Regel für Inferenz-Arbeitslasten verwendet. | xla_all_reduce_latency_bound_threshold_in_bytes=-1 (nicht aktiviert) |
4~16Mb(i.e. 4~16 * 1024 * 1024) |
[0, 9223372036854775807] |
Latenzgebunden xla_collective_permute_latency_bound_threshold_in_bytes |
Dieses Flag ist für latenzgebundene (d.h. kleine) All-Gather-Vorgänge vorgesehen. Wenn Sie diese Option aktivieren, werden bestimmte Optimierungen ausgelöst, die die Ausführungszeit für latenzgebundene kollektive Permutationen verkürzen können. Sie wird in der Regel für Inferenz-Arbeitslasten verwendet. | xla_collective_permute_latency_bound_threshold_in_bytes=-1 (nicht aktiviert) |
4~16Mb(i.e. 4~16 * 1024 * 1024) |
[0, 9223372036854775807] |
Latenzgebunden xla_all_to_all_latency_bound_threshold_in_bytes |
Dieses Flag ist für latenzgebundene (d.h. kleine) All-Gather-Vorgänge vorgesehen. Wenn Sie diese Option aktivieren, werden bestimmte Optimierungen ausgelöst, die die Ausführungszeit für latenzgebundene All-to-All-Vorgänge verkürzen können. Sie wird in der Regel für Inferenz-Arbeitslasten verwendet. | xla_all_to_all_latency_bound_threshold_in_bytes=-1 (nicht aktiviert) |
4~16Mb(i.e. 4~16 * 1024 * 1024) |
[0, 9223372036854775807] |
xla_enable_async_collective_permute |
Schreibt alle kollektiven Permutationsvorgänge in ihre asynchronen Varianten um. Wenn der Wert auf auto festgelegt ist, kann XLA basierend auf anderen Konfigurationen oder Bedingungen automatisch asynchrone Kollektive aktivieren. |
xla_enable_async_collective_permute=kAuto |
xla_enable_async_collective_permute=kAuto |
xla_enable_async_collective_permute=kAuto/kEnabled/kDisabled |
Memory-Flags
Die unten aufgeführten Flags sind für die Behebung von HBM-bezogenen Problemen vorgesehen. Diese sollten nur angepasst werden, wenn bei der Modellkompilierung HBM-Fehler aufgrund fehlenden Speichers auftreten. In allen anderen Fällen werden die Standardwerte empfohlen, da sich Änderungen negativ auf die Leistung auswirken können.
| Flag | Beschreibung | Standardwerte | Vorgeschlagene Werte | Kandidatenwerte |
|---|---|---|---|---|
Planer xla_latency_hiding_scheduler_rerun |
Mit dieser Einstellung wird das Verhalten des Planers zum Verbergen von Latenz angepasst. Dabei wird das für die Planung zugewiesene Speicherlimit mit jedem „Rerun“ des Prozesses schrittweise reduziert. | 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 |
Dieses Flag aktiviert Fusions vom Typ „reduce+broadcast“ und kann die Arbeitsspeichernutzung verringern. | xla_tpu_rwb_fusion=true |
xla_tpu_rwb_fusion=false |
xla_tpu_rwb_fusion=true/false |
Planer xla_memory_scheduler |
Dieses Flag gibt den Algorithmus an, den der Memory-Scheduler verwendet, um den Speicherverbrauch zu minimieren. Mit einem komplexeren Algorithmus kann ein weniger speicherintensiver Zeitplan erstellt werden, allerdings auf Kosten einer längeren Kompilierungszeit. | xla_memory_scheduler=kDefault |
xla_memory_scheduler=kBrkga |
xla_memory_scheduler=kDefault/kList/kDfs/kPostOrder/kBrkga |
Planer xla_tpu_enable_latency_hiding_scheduler |
Dieses Flag aktiviert den Scheduler für die Latenzverdeckung, der es uns ermöglicht, asynchrone statt synchroner Kollektive auszuführen. Wenn Sie diese Funktion deaktivieren, wird der Speicherverbrauch reduziert, allerdings geht die Leistungssteigerung durch diese asynchronen Vorgänge verloren. | 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 |
Mit diesem Flag wird der untere Schwellenwert für die Mindestgröße des Punkts festgelegt, um die kollektive Matmul-Operation auszulösen. Wenn Sie einen höheren Wert festlegen, sparen Sie Arbeitsspeicher, verlieren aber Möglichkeiten, kollektive Matmul-Operationen auszuführen. | xla_jf_spmd_threshold_for_windowed_einsum_mib=-1 |
10Mb~1Gb (i.e. 10*1024*1024 ~ 1024*1024*1024) |
[0, 9223372036854775807] |
Weitere häufig verwendete Flags
| Flag | Typ | Hinweise |
|---|---|---|
xla_dump_to |
String (Dateipfad) | Der Ordner, in dem HLO-Dateien und andere Artefakte vor der Optimierung platziert werden (siehe XLA-Tools). |
TPU XLA-Flags
| Flag | Typ | Hinweise |
|---|---|---|
xla_tpu_enable_data_parallel_all_reduce_opt |
Boolesch (true/false) | Optimierung zur Erhöhung der Überschneidungsmöglichkeiten für All-Reduce-Vorgänge im DCN (Data Center Networking), die für die datenparallele Aufteilung verwendet werden. |
xla_tpu_data_parallel_opt_different_sized_ops |
Boolesch (true/false) | Ermöglicht das Pipelining von datenparallelen Vorgängen über mehrere Iterationen hinweg, auch wenn die Ausgabegrößen nicht mit den Größen übereinstimmen, die in den gestapelten Variablen gespeichert werden können. Kann zu Arbeitsspeichermangel führen. |
xla_tpu_spmd_rng_bit_generator_unsafe |
Boolesch (true/false) | Gibt an, ob der RngBitGenerator-HLO partitioniert ausgeführt werden soll. Dies ist unsicher, wenn deterministische Ergebnisse bei unterschiedlichen Shardings für verschiedene Teile der Berechnung erwartet werden. |
xla_tpu_megacore_fusion_allow_ags |
Boolesch (true/false) | Ermöglicht das Zusammenführen von „all-gather“-Vorgängen mit Faltungen/„all-reduce“-Vorgängen. |
xla_tpu_enable_ag_backward_pipelining |
Boolesch (true/false) | Pipelines werden rückwärts durch Scan-Schleifen zusammengeführt (derzeit Megascale-Zusammenführungen). |
GPU-XLA-Flags
| Flag | Typ | Hinweise |
|---|---|---|
xla_gpu_enable_latency_hiding_scheduler |
Boolesch (true/false) | Dieses Flag ermöglicht es Schedulern, die Latenz verbergen, asynchrone Kommunikation effizient mit Berechnungen zu überlappen. Der Standardwert lautet False. |
xla_gpu_enable_triton_gemm |
Boolesch (true/false) | Triton-basierte Matrixmultiplikation verwenden. |
xla_gpu_graph_level |
Flag (0–3) | Das Legacy-Flag zum Festlegen der GPU-Grafikebene. Verwenden Sie xla_gpu_enable_command_buffer in neuen Anwendungsfällen. 0 = aus; 1 = Fusions und Memcpys erfassen; 2 = GEMMs erfassen; 3 = Faltungen erfassen. |
xla_gpu_all_reduce_combine_threshold_bytes |
Ganzzahl (Byte) | Mit diesen Flags wird festgelegt, wann mehrere kleine AllGather-, ReduceScatter- oder AllReduce-Vorgänge in einen großen AllGather-, ReduceScatter- oder AllReduce-Vorgang zusammengefasst werden, um die für die geräteübergreifende Kommunikation benötigte Zeit zu verkürzen. Bei den AllGather-/ReduceScatter-Schwellenwerten für eine Transformer-basierte Arbeitslast sollten Sie sie beispielsweise so hoch einstellen, dass mindestens das Gewicht von AllGather / ReduceScatter einer Transformer-Schicht kombiniert wird. Standardmäßig ist „combine_threshold_bytes“ auf 256 festgelegt. |
xla_gpu_all_gather_combine_threshold_bytes |
Ganzzahl (Byte) | Siehe oben unter „xla_gpu_all_reduce_combine_threshold_bytes“. |
xla_gpu_reduce_scatter_combine_threshold_bytes |
Ganzzahl (Byte) | Siehe oben unter „xla_gpu_all_reduce_combine_threshold_bytes“. |
xla_gpu_enable_pipelined_all_gather |
Boolesch (true/false) | Aktivieren Sie das Pipelining von All-Gather-Anweisungen. |
xla_gpu_enable_pipelined_reduce_scatter |
Boolesch (true/false) | Pipelining von Reduce-Scatter-Anweisungen aktivieren. |
xla_gpu_enable_pipelined_all_reduce |
Boolesch (true/false) | Aktivieren Sie das Pipelining von All-Reduce-Anweisungen. |
xla_gpu_enable_while_loop_double_buffering |
Boolesch (true/false) | Aktivieren Sie die doppelte Pufferung für die while-Schleife. |
xla_gpu_enable_all_gather_combine_by_dim |
Boolesch (true/false) | Kombinieren Sie All-Gather-Vorgänge mit derselben Gather-Dimension oder unabhängig von ihrer Dimension. |
xla_gpu_enable_reduce_scatter_combine_by_dim |
Boolesch (true/false) | Kombinieren Sie Reduce-Scatter-Vorgänge mit derselben Dimension oder unabhängig von ihrer Dimension. |