facebookresearch/HolisticTraceAnalysis

全面追踪分析 (HTA)

是什么 – HTA 是一个 Python 库，帮助研究人员和工程师理解分布式 PyTorch 训练任务为何变慢。它读取 PyTorch Profiler（Kineto）生成的追踪文件，并将其转化为一系列可直接使用的数据框和可视化图表，精准定位 GPU 在计算、通信、内存操作或空闲状态上花费的时间。

核心功能

功能	你将获得
时间分解	按 GPU 分解时间，分为计算、通信、内存操作和空闲时段。
内核分解与分布	列出每个 rank 上运行时间最长的 CUDA 内核，并展示其在不同 rank 间的持续时间差异。
空闲时间分析	显示 GPU 处于空闲状态的原因（等待主机、等待其他内核、未知）。
通信-计算重叠	通信与有用计算重叠的时间占比。
频繁内核模式	为特定 PyTorch 操作符（如 aten::linear）找出最常重复出现的 CUDA 内核。
内核启动统计	非常短、非常长、或异常延迟的内核启动时间的直方图。
增强计数器	在原始性能分析中添加队列长度和内存带宽追踪，用于类似 Roofline 的分析。
追踪对比	两个运行结果的并排差异分析，查看变化。
CUPTI 计数器 API（实验性）	提取低级别 GPU 性能计数器，支持更深入的 Roofline 分析。

适用人群 – 任何在多个 GPU/主机上训练大型模型，希望以系统化、可脚本化方式定位瓶颈，而无需手动挖掘原始性能分析输出的人。

安装

# 从 PyPI（推荐）
pip install HolisticTraceAnalysis

# 或从源码安装
git clone https://github.com/facebookresearch/HolisticTraceAnalysis.git
cd HolisticTraceAnalysis
git submodule update --init
pip install -r requirements.txt
pip install -e .

支持 Linux 和 macOS，需 Python ≥ 3.10。

典型工作流程

1. 收集追踪数据 使用 PyTorch Profiler（torch.profiler.profile(..., record_shapes=True, with_stack=True, ...)）。
2. 将所有追踪文件放入一个文件夹 – HTA 期望每个运行对应一个独立目录。
3. 打开一个笔记本（或 Python 脚本）并创建 TraceAnalysis 对象：

   from hta.trace_analysis import TraceAnalysis
   analyzer = TraceAnalysis(trace_dir="/path/to/traces")
   

4. 调用所需方法，例如：

   temporal = analyzer.get_temporal_breakdown()
   kernels  = analyzer.get_gpu_kernel_breakdown()
   idle     = analyzer.get_idle_time_breakdown()
   overlap  = analyzer.get_comm_comp_overlap()
   

   每个方法返回一个 pandas DataFrame，可用于绘图或检查。
5. 深入分析时，可使用 examples/ 中的演示笔记本（trace_analysis_demo.ipynb, trace_diff_demo.ipynb）。

配置 – 日志控制通过 hta/configs/logging.config 进行；可更改日志级别或指定自定义配置文件。

文档与支持 – 完整 API 文档托管于 https://hta.readthedocs.io。项目采用 MIT 许可证开源，欢迎通过拉取请求贡献代码。

---

总结 – HTA 将原始 PyTorch Profiler 追踪数据转化为可操作的性能报告，使优化分布式训练流水线变得更加容易。