调查 JAX 追踪
概述
对一次捕获进行七步调查：结论、对账、泳道图、内核 → HLO → Python 代码行、屋顶线、应该改变什么，以及信任列表。
所有信息均来自用户提供的捕获：<host>.xplane.pb 和/或 <host>.trace.json.gz。不要重新运行工作负载、转储 XLA 或运行其他分析器。如果追踪无法回答，则将其列入“不可信”列表。
目标
报告完成后，读者将了解：
- (a) 应该改变什么以使步骤更快，以及估计及其假设；
- (b) 此追踪无法支持归因的地方，并附有计数。
工具
脚本路径相对于此技能的基础目录，即技能加载时显示的目录。从临时目录运行：两个脚本都将输出写入当前目录。
B=<此技能的基础目录>
python3 -P $B/../analyzing-jax-launch-overhead/scripts/skip_stats.py <捕获目录>/<host>.trace.json.gz
python3 -P $B/scripts/attribute_kernels.py <捕获目录> [--program ID] [--user-re REGEX]
skip_stats.py（仅使用标准库）：结论、每个流的时间预算、空闲期间的主机事件以及 SKIP 论文指标。如何阅读：请参考 analyzing-jax-launch-overhead 技能。attribute_kernels.py（需要jax和xprof包）：每个 GPU 内核 → HLO 指令 → op_name → Python 代码行。它只读取.xplane.pb，首先将其复制到./attribution/xplane/，因为 XProf 会将 HLO protos 写入其旁边。它将kernels.csv、ops.csv、lines.csv、how.csv、summary.txt以及程序的 HLO 文本写入./attribution/。- exact（精确）： 追踪器自身的
hlo_op，适用于每个调用启动一个内核的内核。 - name（名称）： 图形内核，其名称是其窗口中的唯一候选者。
- order（顺序）： 共享内核、库 GEMM、复制和集合操作，按调度顺序匹配。
- none（无）： 没有剩余候选者。
- 窗口是程序中相邻精确内核之间的指令，while 循环按其已知的迭代次数展开。
- 验证： 在 MiniGPT 捕获上，它为所有 147,840 个内核选择了与管道相同的指令，该管道与同一程序的命令缓冲区关闭运行 100% 一致。
- exact（精确）： 追踪器自身的
五种技术，每种都有其工具
| 领域 | 技术 | 此处工具 |
|---|---|---|
| 性能工程 | 关键路径；屋顶线 | skip_stats.py 的流预算；仅使用正确峰值的屋顶线（请参阅陷阱） |
| 分布式追踪 | 通过相关 ID 连接主机和 GPU 事件；找到掉队者 | skip_stats.py → rows.csv；跨 GPU 的步骤结束时间分布 |
| 编译器 | 通过编译器传递的源代码来源 | attribute_kernels.py → lines.csv |
| 法证会计 | 总数必须在独立记录间相符 | 第 2 步中的检查 |
| 信息设计 | 对齐的泳道图，每个 GPU 流一个 | 一个步骤，每个 GPU × 流，放在一个时间轴上 |
步骤
方案包含以下七个步骤，全部针对给定的文件运行。只有新捕获才能回答的问题（主机调用的内部作用、修复的效果、批量扫描）将放入第 7 步的“需要新捕获”列表，是否捕获由用户决定。
- 结论和主机端原因： 运行
skip_stats.py。 - 对账。 每个检查要么成立，要么列入信任列表：
- 每个设备的 GPU 事件：json = xplane；
- 每个 GPU 事件都关联一个主机启动调用；
- 每个步骤的内核数量恒定；
- 所有 GPU 执行相同的工作（每个步骤的内核计数相等，或差异已解释）；
- 对于每个流，计算 + 集合 + memcpy + 空闲比例之和等于跨度。
- 泳道图。 对于中位数步骤，将所有 GPU × 流放在一个时间轴上，再加上启动线程的主机事件。这将显示关键路径和掉队者：哪个 GPU 最后完成，以及该 GPU 在等待什么。
- 内核 → HLO → Python 代码行： 运行
attribute_kernels.py。- 报告每个匹配类别占 GPU 时间的比例，以及没有候选者的内核数量。
- 使用
lines.csv和ops.csv获取每行和每操作的表格。 --user-re选择哪些帧算作用户代码。默认情况下，它是 site-packages 和 jax 之外的第一帧。
- 屋顶线， 仅适用于可以从其 HLO 形状获取 FLOPs 和字节数的操作：
attribute_kernels.py写入./attribution/的<program>(<id>).hlo.txt。使用操作数据类型的峰值，引用自供应商规格并标记为 read（读取）；追踪无法测量峰值。 - 应该改变什么。 根据原因在关键路径中的占比进行排序。对于每个原因，估计如果它消失后的步骤时间，例如，关键流上的步骤 − 主机空闲时间。将其标记为 inferred（推断），未运行，并将调节参数列为假设供用户测试（
donate_argnums、关闭命令缓冲区、更少或更大的集合操作）。 - 此追踪不可信的地方： 列出每个项目及其计数或比例。然后添加“需要新捕获”列表：每个未解决的问题，以及什么捕获可以回答它。
陷阱（在 ROCm MI355X 捕获上测量）
- XProf 的操作视图遗漏了图形内核。 在 4-GPU 捕获中，
hlo_stats覆盖了 GPU 内核时间的 37.6%。另外 62.4%（命令缓冲区内核）在kernel_stats中显示为操作“XlaModule”。在引用hlo_stats、op_profile或roofline_model之前检查覆盖率。 - XProf 的 gfx950 屋顶线峰值 为 2,517 TFLOP/s 和 3,624 GiB/s，无论数据类型如何。FP32 GEMM 需要 FP32 峰值。
- RCCL/NCCL 内核时间包括等待对等节点的时间。 这不是传输时间。
- 异步集合操作： 它们的
hlo_op（例如all-gather.34）不在追踪中存储的 HLO 中。attribute_kernels.py保留标签，并按顺序从匹配的集合启动指令中获取 Python 代码行。 - ROCm memcpy 行 可能显示
kind:Unknown size:0。 - Python 追踪器 会夸大主机时间。
- 一个程序可能出现两次（两个具有不同 ID 的
.hlo_proto.pb文件）。attribute_kernels.py使用来自 xplane 的 GPU 时间最多的程序 ID；json 没有program_id。
报告
- Markdown： 先写发现。
- 数字： 每个数字都附带其在步骤和 GPU 上的范围，标记为 measured（测量）、read（读取） 或 inferred（推断）。
- 章节： 结论和原因、对账、每个 GPU/流的表格、顶级操作和 Python 代码行、应该改变什么（估计）以及信任列表。
- 事实核查： 在发送之前，让一个新的子代理根据表格检查每个声明。
- HTML 页面： 如果用户需要，writing-trace-report-pages 技能可以根据此报告及其表格构建 HTML 页面。