XLA 编译器标志列表

内容

XLA 编译器标志列表#

简介#

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

XLA：Jax 背后的引擎#

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

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

在 Jax 中配置 XLA：#

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

对于 colab 笔记本

使用 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 标志#

标志	类型	注释
`xla_dump_to`	字符串（文件路径）	放置预优化 HLO 文件和其他构件的文件夹（请参阅XLA 工具）。
`xla_enable_async_collective_permute`	三态标志（true/false/auto）	将所有集体置换操作重写为其异步变体。当设置为`auto`时，XLA 可以根据其他配置或条件自动开启异步集体。
`xla_enable_async_all_gather`	三态标志（true/false/auto）	如果设置为 true，则启用异步全聚合。如果为`auto`，则仅对实现异步全聚合的平台启用。实现（例如 BC-卸载或延续融合）根据其他标志值选择。
`xla_disable_hlo_passes`	字符串（用逗号分隔的传递名称列表）	要禁用的 HLO 传递的逗号分隔列表。这些名称必须与传递名称完全匹配（逗号之间没有空格）。

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_enable_async_collective_fusion`	布尔值（true/false）	启用将异步集体通信与在其 -start 和 -done 指令之间调度的计算操作（输出/循环融合或卷积）融合的传递。
`xla_tpu_enable_async_collective_fusion_fuse_all_gather`	三态标志（true/false/auto）	启用在 AsyncCollectiveFusion 传递中融合全聚合。如果设置为`auto`，则将根据目标启用。
`xla_tpu_enable_async_collective_fusion_multiple_steps`	布尔值（true/false）	启用在 AsyncCollectiveFusion 传递中分多个步骤（融合）继续相同的异步集体。
`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）	允许将全聚合与卷积/全简化融合。
`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 = 捕获融合和内存复制；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）	启用全聚合指令的流水线处理。
`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_triton_softmax_fusion`	布尔值（true/false）	使用基于 Triton 的 Softmax 融合。
`xla_gpu_enable_all_gather_combine_by_dim`	布尔值（true/false）	组合具有相同聚合维度或与维度无关的全聚合操作。
`xla_gpu_enable_reduce_scatter_combine_by_dim`	布尔值（true/false）	组合具有相同维度或与维度无关的简化散射操作。

其他阅读资料

GPU 性能技巧