本指南精選了重要的 XLA 標記，可協助使用者有效瀏覽及運用 XLA 的功能。以下各節詳細說明可能會大幅影響執行階段效能和記憶體用量的旗標。啟用標記後，如果發生當機等問題，建議還原為預設設定，並建立 GitHub 問題。
成效標記
下列旗標有助於提升執行階段效能。 實驗這些設定可能會大幅提升成效。
| 檢舉 | 說明 | 預設值 | 建議值 | 候選值 |
|---|---|---|---|---|
| 管道 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 |
這 3 個標記應一併使用，以啟用 ICI(晶片間互連) 全體收集作業的集體管線化，進而創造更多執行作業重疊的機會。 | 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 |
這 3 個旗標應搭配使用，以在 v5e 上啟用非同步 all-gather 作業。 | 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 |
這 2 個旗標應搭配使用，以在 v5e 上啟用非同步 all-reduce 作業。 | 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 |
這個標記可啟用非同步的 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 |
延遲時間限制 xla_all_gather_latency_bound_threshold_in_bytes |
這個標記適用於延遲時間受限 (即小型) 的 all-gather 作業。啟用這項功能會觸發特定最佳化作業，可縮短延遲時間受限的全體收集作業執行時間。通常用於推論工作負載。 | xla_all_gather_latency_bound_threshold_in_bytes=-1 (未啟用) |
4~16Mb(i.e. 4~16 * 1024 * 1024) |
[0, 9223372036854775807] |
延遲時間限制 xla_all_reduce_latency_bound_threshold_in_bytes |
這個標記適用於延遲時間受限 (即小型) 的 all-gather 作業。啟用這項功能會觸發特定最佳化作業，可縮短延遲時間受限的 all-reduce 執行時間。通常用於推論工作負載。 | xla_all_reduce_latency_bound_threshold_in_bytes=-1 (未啟用) |
4~16Mb(i.e. 4~16 * 1024 * 1024) |
[0, 9223372036854775807] |
延遲時間限制 xla_collective_permute_latency_bound_threshold_in_bytes |
這個標記適用於延遲時間受限 (即小型) 的 all-gather 作業。啟用這項功能會觸發特定最佳化作業，可縮短延遲時間受限的集體排列執行時間。通常用於推論工作負載。 | xla_collective_permute_latency_bound_threshold_in_bytes=-1 (未啟用) |
4~16Mb(i.e. 4~16 * 1024 * 1024) |
[0, 9223372036854775807] |
延遲時間限制 xla_all_to_all_latency_bound_threshold_in_bytes |
這個標記適用於延遲時間受限 (即小型) 的 all-gather 作業。啟用這項功能會觸發特定最佳化作業，可縮短延遲時間受限的 all-to-all 執行時間。通常用於推論工作負載。 | xla_all_to_all_latency_bound_threshold_in_bytes=-1 (未啟用) |
4~16Mb(i.e. 4~16 * 1024 * 1024) |
[0, 9223372036854775807] |
xla_enable_async_collective_permute |
將所有集體排列作業重新編寫為非同步變體。設為 auto 時，XLA 可根據其他設定或條件自動開啟非同步集合。 |
xla_enable_async_collective_permute=kAuto |
xla_enable_async_collective_permute=kAuto |
xla_enable_async_collective_permute=kAuto/kEnabled/kDisabled |
記憶體旗標
以下列出的標記用於解決 HBM 相關問題。只有在模型編譯期間發生 HBM「記憶體不足」錯誤時，才應調整這些值。在所有其他情況下，建議使用預設值，因為變更這些值可能會對成效造成負面影響。
| 檢舉 | 說明 | 預設值 | 建議值 | 候選值 |
|---|---|---|---|---|
排程器 xla_latency_hiding_scheduler_rerun |
這項設定可調整延遲隱藏排程器的行為。這項功能會隨著程序的每次「重新執行」，逐步減少分配給排程的記憶體上限。 | 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 |
這個標記可啟用 reduce+broadcast 類型的融合，並減少記憶體用量。 | xla_tpu_rwb_fusion=true |
xla_tpu_rwb_fusion=false |
xla_tpu_rwb_fusion=true/false |
排程器 xla_memory_scheduler |
這個標記會指定記憶體排程器使用的演算法，以盡量減少記憶體用量。使用更進階的演算法可能會產生耗用較少記憶體的排程，但編譯時間會較長。 | xla_memory_scheduler=kDefault |
xla_memory_scheduler=kBrkga |
xla_memory_scheduler=kDefault/kList/kDfs/kPostOrder/kBrkga |
排程器 xla_tpu_enable_latency_hiding_scheduler |
這個標記會啟用延遲隱藏排程器，讓我們執行非同步集體作業，而非同步作業。停用這項功能可減少記憶體用量，但會失去這些非同步作業帶來的效能提升。 | 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 |
這個旗標會設定點的最小尺寸下限，以觸發集體矩陣乘法。將這個值設為較高，可節省記憶體，但會失去執行集體矩陣乘法的機會。 | xla_jf_spmd_threshold_for_windowed_einsum_mib=-1 |
10Mb~1Gb (i.e. 10*1024*1024 ~ 1024*1024*1024) |
[0, 9223372036854775807] |
其他常用旗標
| 檢舉 | 類型 | 附註 |
|---|---|---|
xla_dump_to |
字串 (檔案路徑) | 預先最佳化的 HLO 檔案和其他構件的放置資料夾 (請參閱 XLA 工具)。 |
TPU XLA 旗標
| 檢舉 | 類型 | 附註 |
|---|---|---|
xla_tpu_enable_data_parallel_all_reduce_opt |
布林值 (true/false) | 最佳化 DCN (資料中心網路) 全縮減，用於資料平行分片，以增加重疊機會。 |
xla_tpu_data_parallel_opt_different_sized_ops |
布林值 (true/false) | 即使資料平行作業的輸出大小與堆疊變數中可儲存的大小不符，也能在多次疊代中進行資料平行作業的管線化。可能會增加記憶體壓力。 |
xla_tpu_spmd_rng_bit_generator_unsafe |
布林值 (true/false) | 是否以分割方式執行 RngBitGenerator HLO。如果預期在運算的各個部分使用不同的分片，並取得決定性結果，則這種做法並不安全。 |
xla_tpu_megacore_fusion_allow_ags |
布林值 (true/false) | 允許將 all-gather 與捲積/all-reduce 融合。 |
xla_tpu_enable_ag_backward_pipelining |
布林值 (true/false) | 管道會透過掃描迴圈，向後全數收集 (目前為巨量全數收集)。 |
GPU XLA 旗標
| 檢舉 | 類型 | 附註 |
|---|---|---|
xla_gpu_enable_latency_hiding_scheduler |
布林值 (true/false) | 這個標記可讓延遲隱藏排程器有效率地重疊非同步通訊與運算。預設值為 False。 |
xla_gpu_enable_triton_gemm |
布林值 (true/false) | 使用以 Triton 為基礎的矩陣乘法。 |
xla_gpu_graph_level |
旗標 (0-3) | 用於設定 GPU 圖形層級的舊版標記。在新用途中，使用 xla_gpu_enable_command_buffer。0 = 關閉；1 = 擷取融合和 memcpys；2 = 擷取 gemms；3 = 擷取捲積。 |
xla_gpu_all_reduce_combine_threshold_bytes |
整數 (位元組) | 這些旗標會調整何時將多個小型 AllGather / ReduceScatter / AllReduce 合併為一個大型 AllGather / ReduceScatter / AllReduce，以減少跨裝置通訊所花費的時間。舉例來說，如果是以 Transformer 為基礎的工作負載，請考慮將 AllGather / ReduceScatter 門檻調高，至少要合併 Transformer 層的權重 AllGather / ReduceScatter。根據預設，combine_threshold_bytes 會設為 256。 |
xla_gpu_all_gather_combine_threshold_bytes |
整數 (位元組) | 請參閱上方的 xla_gpu_all_reduce_combine_threshold_bytes。 |
xla_gpu_reduce_scatter_combine_threshold_bytes |
整數 (位元組) | 請參閱上方的 xla_gpu_all_reduce_combine_threshold_bytes。 |
xla_gpu_enable_pipelined_all_gather |
布林值 (true/false) | 啟用 all-gather 指令的管線化。 |
xla_gpu_enable_pipelined_reduce_scatter |
布林值 (true/false) | 啟用減少分散指令的管線化。 |
xla_gpu_enable_pipelined_all_reduce |
布林值 (true/false) | 啟用所有縮減指令的管線化。 |
xla_gpu_enable_while_loop_double_buffering |
布林值 (true/false) | 為 while 迴圈啟用雙重緩衝。 |
xla_gpu_enable_all_gather_combine_by_dim |
布林值 (true/false) | 合併所有收集作業，無論收集維度是否相同。 |
xla_gpu_enable_reduce_scatter_combine_by_dim |
布林值 (true/false) | 將 reduce-scatter 作業與相同維度或不論維度合併。 |