SuperAPI AI 工具集
合作ENAI 行业资讯收藏
Skill 广场cuda_graph_trace_compaction

cuda_graph_trace_compaction

数据与分析

查看项目

更新于 10-22 04:51

技能介绍

用途

CUDA graph replay 会为每个捕获区域分配独立 stream，导致 PyTorch Profiler 追踪在 Perfetto 中展开为数百条近乎空行的行，难以阅读。本工具通过 scripts/compact_cuda_graph_trace.py 将切片重新打包到少数语义泳道（compute、NCCL all-gather、reduce-scatter、……、memcpy、memset），并可选择合并按 rank 划分的流水线并行追踪。

主要能力

  • 单一追踪压缩：将一个 rank 的 trace 重新分组为语义泳道。
  • PP 合并：将目录中所有 rank*_trace.json[.gz] 合并为 pp_traces_merged_compacted.json.gz，按 pipeline 顺序堆叠各 rank，并在 torch.distributed.pipelining 发出的 PP:<stage><op><microbatch> 注释之间绘制流箭头（跨 rank 的 SEND_F -> RECV_F，同一 rank 内的 RECV_F -> F -> SEND_F）。
  • 单 rank 压缩包含 rank 本地 PP 箭头：UNSHARD -> compute、RECV -> compute、compute -> SEND。
  • 合并时还会把前两条 all-gather 泳道上每个 PP:<stage>UNSHARD 注释连接到第一条 compute 泳道上其后第一个同 stage 的 compute 注释。
  • 无注释时不绘制箭头，但合并仍可正常工作。
  • 合并时按共同 baseTimeNanoseconds 对齐各 rank，并按 pipeline 顺序堆叠。
  • 脚本会打印 JSON 摘要并自校验：若实测切片发生变化或某泳道出现重叠切片则报错。
  • 摘要中 original_streams 与 compacted_lanes 可用于查看压缩比；pp_send_recv_unmatched 非零通常意味着捕获窗口被截断，而非流水线缺陷。

使用方式

# 单追踪 -> 输入旁生成 <name>_compacted.json.gz
python3 .claude/skills/cuda_graph_trace_compaction/scripts/compact_cuda_graph_trace.py \
    outputs/<run>/profile_trace/iteration_10/rank0_trace.json.gz

# 目录中所有 rank*_trace.json[.gz] -> pp_traces_merged_compacted.json.gz
python3 .claude/skills/cuda_graph_trace_compaction/scripts/compact_cuda_graph_trace.py \
    outputs/<run>/profile_trace/iteration_10 --merge-pp-ranks

-o 用于设置输出路径，--overwrite 用于替换已有输出。rank 顺序取自 distributedInfo.pg_config 中的 torchtitan_real_pp 组或 mesh_pp 组描述；若不存在，则按全局 rank 排序，这仅在 PP rank 顺序与全局 rank 顺序一致时才正确。

适用场景

适合需要在 Perfetto 中阅读 CUDA graph 密集追踪、并理解流水线并行 SEND/RECV 与计算关系的场景。

限制

  • 仅用于可视化：切片保留原名、pid、ts、dur 与 args，仅改变显示位置（tid/args.stream）。测量应使用原始文件，阅读使用压缩副本。
  • 脚本由 agent 维护，不是 torchtitan API；应针对手头追踪直接修改，而不是绕开它。常见调整点包括 _classify_slice 与 GROUP_ORDER/GROUP_LABELS（新增泳道分组）、PP_ANNOTATION_PATTERN（不同注释方案）。修改后需保持测试通过。
  • 测试命令：pytest .claude/skills/cuda_graph_trace_compaction/scripts/，使用手工构造的追踪，无需 GPU。该目录在 testpaths 之外，因此 pytest tests/ 会跳过它。修改脚本后应运行测试；若不变式被破坏，会产出看似合理但错误的追踪而非报错。