企业官网建设流程全解析

TensorFlow 如何用 XLA_FLAGS 环境变量 dump HLO 文本辅助调试？

【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow

当你的 TensorFlow 程序通过 XLA 编译后出现结果异常、编译失败或想向 XLA 提交可复现的 bug report 时，需要拿到程序实际生成的 HLO（XLA 的中间表示，列出 HLO 指令及其属性的人类可读文本）。XLA 官方文档明确指出：HLO dump 是调试 XLA 问题和撰写 bug report 时的常用产物。在 TensorFlow 场景下，最常用的方式是通过XLA_FLAGS环境变量指定--xla_dump_to把 HLO 模块 dump 到指定目录，再结合TF_XLA_FLAGS把聚类与 cluster 嵌入图也一并导出。适用前提：你的程序已经（或即将）以 XLA 方式编译运行，即使用了@tf.function(jit_compile=True)、Keras 的model.compile(jit_compile=True)，或者通过TF_XLA_FLAGS=--tf_xla_auto_jit=2开启 auto-clustering。

先让程序跑在 XLA 上：两种启用方式

HLO dump 只在 XLA 编译发生时产生，所以第一步是确认编译路径。

方式一：显式编译。用jit_compile=True精确控制哪些函数交给 XLA：

@tf.function(jit_compile=True) def train_mnist(images, labels): # 训练步骤 ...

Keras 模型则传给model.compile：

model.compile(optimizer="adam", jit_compile=True)

注意两点限制，都来自 tf2xla 文档：

• jit_compile是 must-compile 语义：整个函数要么整体被 XLA 编译，要么抛出errors.InvalidArgumentError。例如包含tf.unique(x)的函数会编译失败，因为无法在不运行整个计算的情况下推断所有张量维度；
• 嵌套函数中，只要调用栈里至少一个函数是jit_compile=True，该栈上的函数都会被编译。

方式二：auto-clustering（不改代码）。通过TF_XLA_FLAGS环境变量自动找出可编译的子图（cluster）交给 XLA：

TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program

auto-clustering 目前针对 GPU 优化；CPU 上可用TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit"额外开启，但 CPU 与多 GPU 环境下的 auto-clustering 是实验性的。TF2 中只有tf.function内部的代码会被聚类。path/to/your/tf/program替换为你自己的 TensorFlow 程序入口。

用 XLA_FLAGS 把 HLO 文本 dump 到目录

在程序启动前设置XLA_FLAGS环境变量（JAX、TensorFlow、PyTorch/XLA 通用）。--xla_dump_to后面是 dump 目录，文档示例使用/tmp或/tmp/xladump，目录需可写：

XLA_FLAGS="--xla_dump_to=DIRECTORY_PATH"

默认情况下，这会在优化流水线的最开始和最后各 dump 一次 HLO 模块文本。也可以显式指定格式，以下变体都要求同时给出--xla_dump_to：

# 文本 dump XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=DIRECTORY_PATH" # HLO proto（机器可读的结构化格式） XLA_FLAGS="--xla_dump_hlo_as_proto --xla_dump_to=DIRECTORY_PATH" # HLO Snapshot（HLO 模块加输入数据，用于回放真实输入而非随机数据） XLA_FLAGS="--xla_dump_hlo_snapshots --xla_dump_to=DIRECTORY_PATH" # graphviz URL / HTML 图渲染，只适合小图 XLA_FLAGS="--xla_dump_hlo_as_url --xla_dump_to=DIRECTORY_PATH" XLA_FLAGS="--xla_dump_hlo_as_html --xla_dump_to=DIRECTORY_PATH"

两点行为规则（来自 hlo_dumps.md）：

• 如果没指定--xla_dump_to但设置了其他 dump flag，dump 会输出到 stdout；但二进制数据（如 proto 文件）不会写到 stdout；
• 大图不建议用 graphviz 渲染，文档建议用interactive_graphviz可视化图的一部分。

调试时如果只想看特定优化 pass 之后的 HLO，而不是只看流水线首尾两个阶段，用正则匹配 pass 名：

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

例如 dump SPMD 分区相关 pass 的结果：

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

--xla_dump_hlo_pass_re=.*会 dump 每个 pass 之后的结果，文件非常多，谨慎使用。

另一个可选 flag 是--xla_syntax_sugar_async_ops=true：让 dump 使用语法糖包装器作为 op 名，文档称可减少约 20% 的 dump 体积，默认值为false（dump 中使用真实 op 名）。

TensorFlow 场景的完整命令：同时导出聚类信息

对于跑 auto-clustering 的 TensorFlow 程序，tf2xla 文档 给出的组合命令是同时导出 HLO 文本和 XLA cluster 在 TF 图中的嵌入方式：

TF_DUMP_GRAPH_PREFIX=/tmp/generated \ TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \ XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \ my/tensorflow/program

其中my/tensorflow/program替换为你的 TensorFlow 程序入口。程序运行结束后，/tmp/generated目录里会出现：

• module_XXXX.*_optimizations.txt：生成的 XLA 程序文本（HLO），每个编译 cluster 对应一个，提交 XLA bug report 时附上这些文件非常有帮助；
• module_XXXX.ir-*.ll：LLVM 中间表示（含 NVPTX intrinsics）；
• module_XXXX.ptx：生成的 PTX 文件。

如果只想看 XLA cluster 如何嵌入 TF 图（不 dump HLO），文档给的更轻量的组合是：

TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"

从一堆文件里找到出问题的计算

一次运行通常会 dump 出多个计算的文件。文件名里带有 JAX、TensorFlow 或 PyTorch/XLA 日志中出现的 "computation name"，可以直接对应日志定位相关模块，例如：

1624325116260738.module_0065.pmap__unnamed_wrapped_function_.186875.before_optimizations.txt

对不上名字时，可以用 ripgrep 在 dump 目录里按符号搜索，快速判断哪个 module 包含出问题的计算。文档还建议在 bug report 中附上感兴趣的 before/after/buffer-assignment 共 3 个 dump 文件。

对 dump 出来的 HLO 做进一步操作

拿到.txt格式的 HLO 之后，后续处理都有对应工具（详见 XLA Tooling 文档）：

• replay（回放）：用run_hlo_module在指定后端上重放 dump 的计算，可用假数据或 HLO Snapshot 中的真实输入，方便迭代复现问题。CPU / GPU 后端示例：

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

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

• 格式转换：hlo-opt可以在 HLOProto 与文本格式之间互转，适合"只拿到 proto、需要文本"或反过来的情况：

hlo-opt --emit-proto input.hlo # HLO Text -> HLO Proto hlo-opt input.pbtxt # HLO Proto -> HLO Text

• 隔离到单个 XLA 程序：文档建议尽量用run_hlo_module对生成出的单个 XLA 程序反复运行，把 bug 隔离到最小的计算单元。

边界与限制

• HLO dump 依赖 XLA 编译实际发生：没用jit_compile=True、Kerasjit_compile=True或TF_XLA_FLAGSauto-clustering 的程序不会产生这些产物；
• jit_compile=True无法编译维度不可推断的函数（如tf.unique），此时要么改用 auto-clustering 让可编译部分被聚类，要么先调整函数写法；
• graphviz/HTML 图渲染只适合小图；--xla_dump_hlo_pass_re=.*会生成大量文件；
• CPU 与多 GPU 上的 auto-clustering 是实验性特性。

参考文档：Dump HLO Computations、XLA for TensorFlow、XLA Tooling。

【免费下载链接】tensorflowAn Open Source Machine Learning Framework for Everyone项目地址: https://gitcode.com/GitHub_Trending/te/tensorflow