XLA 编译器标志列表

XLA 编译器标志列表#

简介#

本指南简要概述了 XLA 以及 XLA 与 Jax 的关系。有关深入的详细信息，请参阅 XLA 文档。然后列出了常用的 XLA 编译器标志，旨在优化 Jax 程序的性能。

XLA：Jax 背后的动力#

XLA（加速线性代数）是用于线性代数的特定领域编译器，在 Jax 的性能和灵活性方面起着关键作用。它使 Jax 能够通过将您的 Python/NumPy 类代码转换和编译为高效的机器指令，为各种硬件后端（CPU、GPU、TPU）生成优化的代码。

Jax 使用 XLA 的 JIT 编译功能，在运行时将您的 Python 函数转换为优化的 XLA 计算。

在 Jax 中配置 XLA：#

您可以通过在运行 Python 脚本或 colab notebook 之前设置 XLA_FLAGS 环境变量来影响 XLA 在 Jax 中的行为。

对于 colab notebook

使用 os.environ['XLA_FLAGS'] 提供标志

import os

# Set multiple flags separated by spaces
os.environ['XLA_FLAGS'] = '--flag1=value1 --flag2=value2'

对于 python 脚本

将 XLA_FLAGS 指定为 cli 命令的一部分

XLA_FLAGS='--flag1=value1 --flag2=value2'  python3 source.py

重要提示

在导入 Jax 或其他相关库之前设置 XLA_FLAGS。在后端初始化后更改 XLA_FLAGS 将不起作用，并且由于给定的后端初始化时间未明确定义，因此通常在执行任何 Jax 代码之前设置 XLA_FLAGS 更安全。
尝试使用不同的标志来优化您的特定用例的性能。

更多信息

有关 XLA 的完整和最新的文档，请参见官方 XLA 文档。
对于开源 XLA 版本（CPU、GPU）支持的后端，XLA 标志及其默认值在 xla/debug_options_flags.cc 中定义，完整的标志列表可以在此处找到。
TPU 编译器标志不是 OpenXLA 的一部分，但常用选项在下面列出。
请注意，此标志列表并非详尽无遗，并且可能会发生更改。这些标志是实现细节，并且不保证它们将保持可用或保持其当前行为。

常用 XLA 标志#

标志	类型	Notes
`xla_dump_to`	字符串 (文件路径)	将放置预优化 HLO 文件和其他工件的文件夹（请参阅 XLA 工具）。
`xla_enable_async_collective_permute`	TristateFlag (true/false/auto)	将所有 collective-permute 操作重写为其异步变体。当设置为 `auto` 时，XLA 可以根据其他配置或条件自动启用异步 collective。
`xla_enable_async_all_gather`	TristateFlag (true/false/auto)	如果设置为 true，则启用异步 all gather。如果 `auto`，则仅对实现异步 all-gather 的平台启用。实现方式（例如 BC 卸载或连续融合）根据其他标志值选择。
`xla_disable_hlo_passes`	字符串（逗号分隔的 pass 名称列表）	要禁用的 HLO pass 的逗号分隔列表。这些名称必须与 pass 名称完全匹配（逗号周围没有空格）。

TPU XLA 标志#

标志	类型	Notes
`xla_tpu_enable_data_parallel_all_reduce_opt`	布尔值 (true/false)	优化以增加用于数据并行分片的 DCN（数据中心网络）all-reduce 的重叠机会。
`xla_tpu_data_parallel_opt_different_sized_ops`	布尔值 (true/false)	即使数据并行操作的输出大小与堆叠变量中可以就地保存的大小不匹配，也允许跨多个迭代流水线化数据并行操作。可能会增加内存压力。
`xla_tpu_enable_async_collective_fusion`	布尔值 (true/false)	启用 pass，该 pass 将异步集体通信与计划在其 -start 和 -done 指令之间的计算操作（输出/循环融合或卷积）融合。
`xla_tpu_enable_async_collective_fusion_fuse_all_gather`	TristateFlag (true/false/auto)	启用在 AsyncCollectiveFusion pass 中融合 all-gather。如果设置为 `auto`，则将根据目标启用它。
`xla_tpu_enable_async_collective_fusion_multiple_steps`	布尔值 (true/false)	允许在 AsyncCollectiveFusion pass 的多个步骤（融合）中继续相同的异步 collective。
`xla_tpu_overlap_compute_collective_tc`	布尔值 (true/false)	在单个 TensorCore 上启用计算和通信的重叠，即一个核心等效于 MegaCore 融合。
`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)	通过扫描循环向后流水线化 all-gather（当前为 megascale all-gather）。

GPU XLA 标志#

标志	类型	Notes
`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 = 捕获融合和 memcopy；2 = 捕获 gemm；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)	启用 reduce-scatter 指令的流水线化。
`xla_gpu_enable_pipelined_all_reduce`	布尔值 (true/false)	启用 all-reduce 指令的流水线化。
`xla_gpu_enable_while_loop_double_buffering`	布尔值 (true/false)	为 while 循环启用双缓冲。
`xla_gpu_enable_all_gather_combine_by_dim`	布尔值 (true/false)	合并具有相同 gather 维度的 all-gather 操作，或不考虑其维度。
`xla_gpu_enable_reduce_scatter_combine_by_dim`	布尔值 (true/false)	合并具有相同维度的 reduce-scatter 操作，或不考虑其维度。

拓展阅读

GPU 性能技巧