转储 HLO 计算

HLO 转储是计算不同阶段的 HLO 模块的文本表示形式。它对于调试很有用，并且您通常需要在 bug 报告中包含它。这通常是一个可读的文本文件，其中列出了 HLO 指令和它们的属性。有时，HLO 模块会转储为：

  • HloProto：协议缓冲区文件，是一种结构化程度更高、机器可读的格式。
  • HloSnapshot：HLO 模块及其输入。对于重放 HLO，您有时需要提供给指定计算的实际输入，而不是随机数据。

您可以使用 XLA 标志来指定和获取转储。在大多数情况下，您可以使用环境变量来设置它。JAX 还提供了一种以编程方式打印 HLO 转储的方法。

本地执行

使用环境变量

您可以设置 XLA_FLAGS 环境变量，并添加必要的标志来获取内存转储。此方法适用于 JAX、TensorFlow 和 PyTorch/XLA。

如需将 HLO 模块和其他调试信息转储到特定目录，请使用 --xla_dump_to 标志运行程序：

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

例如，您可以使用 /tmp 或 /tmp/xladump 作为路径。

默认情况下，此选项会在优化流水线的开头和结尾将 HLO 模块转储为文本。

您还可以明确指定格式：

  1. 文本转储
XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH"
  1. HLO proto
XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH"
  1. HLO 快照
XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH"
  1. 使用 Graphviz 服务器渲染图（仅适用于小型图）
XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH"
  1. 将图表渲染为 HTML 文件（仅适用于小型图表）
XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

对于较大的图，您可以使用 interactive_graphviz 直观呈现图的部分内容。

转储特定中间传递

除了标准的预优化 / 最终优化 HLO 之外，您还可以转储特定编译器传递后的 HLO 状态。

XLA_FLAGS="--xla_dump_hlo_pass_re=regex --xla_dump_to=DIRECTORY_PATH"

系统将转储名称与正则表达式 (regex) 匹配的传递的 HLO 模块。例如，您可以使用以下命令观察与 SPMD 分区相关的 pass 所产生的 HLO：

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=spmd|propagation"

如需在每次 XLA 传递后转储结果（这会生成大量文件），您可以设置：

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH --xla_dump_hlo_pass_re=.*"

JAX 特有选项

以编程方式在 JAX 中

除了传递标志或环境变量之外，您还可以使用 JAX 的 lower 和 compile API 以编程方式转储 HLO。

使用以下命令在本地提取未优化的原始降级 HLO：

jax.jit(f).lower(*args).as_text('hlo')

如需在 HLO 编译过程中将内容转储到文件，请指定：

compilation_args = {
    'xla_dump_to': DIRECTORY_PATH,
    'xla_dump_hlo_pass_re': 'spmd|propagation', # or some other pass filter
    ...
    }

jax.jit(f).lower(*args).compile(compilation_args)

转储 jaxpr

jaxprs 是 JAX 的程序轨迹中间表示形式。如需转储此信息，请设置以下环境变量：

JAX_DUMP_IR_TO="DIRECTORY_PATH" JAX_DUMP_IR_MODES=jaxpr

如需了解详情，请参阅 JAX 文档中的导出和序列化分阶段计算：调试。

Google Colab

环境变量

在笔记本的第一个执行单元格中（因为环境变量和命令行标志通常只处理一次，例如在模块导入时或 XLA 后端初始化时），添加上述 XLA_FLAGS，并附带 os.environ，例如：

import os
os.environ['XLA_FLAGS'] = "--xla_dump_to=DIRECTORY_PATH"

这会将计算结果转储到 DIRECTORY_PATH，例如 /tmp。在 Colab 中，前往左侧边栏中的“文件”浏览器，查看并访问此目录。

您可以使用“本地执行”部分中提及的所有标志。

JAX 特有选项

与本地执行类似；对于实时交互式自省，您可以直接打印计算的预优化 HLO：

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

c = jax.jit(f).lower(3.).compiler_ir('hlo')

print(c.as_hlo_text())

您还可以直接打印计算的优化 HLO：

def optimized_HLO(f, *args, platform=None):
    print(jax.jit(f).lower(*args).compile().as_text())

def f(x):
    return jax.numpy.sin(jax.numpy.cos(x))

optimized_HLO(f, 1.0)

转储所有/小型计算

如果您想在转储中查看所有内容（包括所有小型编译），请设置 JAX 环境变量：

JAX_COMPILER_DETAILED_LOGGING_MIN_OPS=0

马赛克

Mosaic 是 Pallas TPU 后端和实验性 Pallas GPU 后端的编译器。如需转储镶嵌计算，请设置以下标志：

--xla_mosaic_dump_to=/tmp/mosaic_dumps

或者，将 TPU 初始化实参设置为环境变量：

export LIBTPU_INIT_ARGS="--xla_mosaic_dump_to=/tmp/mosaic_dumps"

如需了解详情，请参阅 JAX 关于 Pallas 和 Mosaic 的文档。

更多 HLO 转储

找到合适的计算方式

通常，许多计算会被转储。转储的文件会明确命名为日志中提及的 JAX、Tensorflow 或 PyTorch/XLA“计算名称”，从而方便您识别相关的 HLO 文件。例如：

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

否则，您可以使用 ripgrep 快速识别哪个模块包含特定符号或计算。

提示：在错误报告中包含您感兴趣的 3 个转储文件（之前/之后/缓冲区分配）。

HLO 转化

一种名为 hlo-opt 的工具，可在 HLOProto 和文本格式之间进行转换。 如果您有一种格式，但需要另一种格式进行调试，此功能会非常有用。

了解如何使用它：XLA 工具文档：hlo-opt。

重放

您可以使用虚假数据或输入快照在指定的 XLA 后端上运行（重放）转储的计算。这是一种方便的方法，可用于重现、迭代和调试 XLA 中的问题。

以下命令使用虚假数据。如果您已保存 HLO 快照，则可以改为传入这些快照，系统将使用快照中的数据。如需在运行快照时仍使用虚假数据，请传递标志 --force_fake_data。

CPU 后端：

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=cpu
 /tmp/xladump/module_4561.before_optimizations.txt

GPU 后端：

bazel run -c opt //xla/hlo/tools:run_hlo_module -- --platform=CUDA
 /tmp/xladump/module_4561.before_optimizations.txt