profile-model

Profile Model

分析 PyTorch profiler 跟踪文件（.json/.json.gz），诊断 GPU 性能问题、查找缓慢的内核并识别瓶颈。

---

Step 1：获取跟踪文件路径

询问用户跟踪文件的路径。跟踪文件通常是：

• 单个文件：/path/to/trace.json 或 /path/to/trace.json.gz
• 目录（多 rank）：/path/to/traces/，其中包含 rank0.json、rank1.json 等文件。

Step 2：将跟踪文件加载为 DataFrame

将跟踪文件加载到 pandas DataFrame，以便灵活查询。

⚠️ 重要：缓存前先询问用户。 不要自动缓存。请询问：

“您是否希望我将此跟踪文件缓存为 Parquet，以便今后更快地加载？ 这会将文件保存到 /tmp/trace_cache/，加快后续加载速度。 （建议用于大型跟踪文件，或您计划提出后续问题的情况）”

只有在用户同意后才进行缓存。

DataFrame 结构：

• 标准列：name、cat、ph、ts、dur、pid、tid、id
• args 展开为：args_device、args_stream、args_Input Dims 等列

加载方法：

1. 解析 JSON/JSON.gz 跟踪文件
2. 将每个事件展开为一行，并将 args 字典展开为 args_* 列
3. 将 ts 和 dur 转换为数值类型

Parquet 缓存注意事项（仅在用户同意时）：

• 缓存到 /tmp/trace_cache/{basename}.parquet
• 通过比较 mtime 与源文件来检查缓存是否为最新
• 处理混合类型：args_* 列通常包含混合类型（int、str、None）。保存为 Parquet 前，检测混合类型列并将其转换为字符串。将 'nan'/'None' 字符串还原为 None。
• 如果 Parquet 保存失败，则将所有 args_* 列转换为字符串后重试
• 如果仍然失败，则继续操作，不使用缓存

Step 3：快速摘要

加载后，显示快速摘要：

• 事件总数及按类别划分的数量（df["cat"].value_counts()）
• 内核数量和内核总耗时
• 按耗时排列的主要内核

Step 4：使用 DataFrame 查询进行分析

用于筛选的关键列值：

• GPU 内核：cat == "kernel"
• CPU 操作：cat == "cpu_op"
• NCCL/通信：name.str.contains("nccl", case=False)
• 内存操作：cat 包含 "mem" 或 "memcpy"

时间戳（ts）和持续时间（dur）的单位是微秒（μs）。

使用标准 pandas 操作（groupby、filter、sort、aggregate）回答用户的问题。

Step 5：解读结果

发现	可能的问题	建议
大量小型内核（平均 < 10μs）	内核启动开销	使用 torch.compile，融合操作
单个内核占主导（>50%）	对计算密集型任务而言属正常情况	检查内核是否经过优化
大量内存操作	内存受限	减少拷贝，使用固定内存
通信耗时高	分布式开销	让通信与计算重叠

Step 6：使用 HTA 进行深入分析（可选）

如果用户希望进行更深入的分析，请使用 HTA（Holistic Trace Analysis）：

# Check if HTA is available
try:
    from hta.trace_analysis import TraceAnalysis
except ImportError:
    print("HTA not installed. Install with: pip install HolisticTraceAnalysis")

HTA 分析函数

用户问题	HTA 函数	返回内容
“GPU 时间都花在哪了？”	get_gpu_kernel_breakdown()	内核耗时分布
“关键路径在哪里？”	critical_path_analysis()	执行过程中的瓶颈路径
“为什么 GPU 空闲？”	get_idle_time_breakdown()	主机等待与内核等待的细分情况
“通信/计算有重叠吗？”	get_comm_comp_overlap()	每个 rank 的重叠百分比
“有没有较慢的 rank？”	get_potential_stragglers()	识别拖后腿的 rank
“总体耗时分布如何？”	get_temporal_breakdown()	空闲/计算/非计算时间

HTA 使用示例

from hta.trace_analysis import TraceAnalysis

# Load traces
analyzer = TraceAnalysis(trace_dir="/path/to/traces/")
# OR: analyzer = TraceAnalysis(trace_files={0: "/path/to/trace.json.gz"})

# GPU Kernel Breakdown
kernel_type_df, kernel_df = analyzer.get_gpu_kernel_breakdown(visualize=False)

# Temporal Breakdown (idle vs compute vs non-compute)
temporal_df = analyzer.get_temporal_breakdown(visualize=False)

# Idle Time Breakdown (why is GPU idle?)
idle_df, _ = analyzer.get_idle_time_breakdown(ranks=[0], visualize=False)

# Critical Path Analysis (for distributed training)
cp_graph, success = analyzer.critical_path_analysis(
    rank=0, annotation="ProfilerStep", instance_id=1
)
if success:
    analyzer.overlay_critical_path_analysis(
        rank=0, critical_path_graph=cp_graph, output_dir="/tmp/cp_output"
    )

# Communication-Computation Overlap
overlap_df = analyzer.get_comm_comp_overlap(visualize=False)

# Straggler Detection
stragglers = analyzer.get_potential_stragglers(num_candidates=2)

---

常见性能问题及修复方法

症状	诊断	修复方法
大量小型内核	内核启动开销	使用 torch.compile、算子融合
空闲时间过长（等待主机）	CPU 瓶颈	增加 DataLoader 的 num_workers，使用 pin_memory=True
空闲时间过长（等待内核）	内核连续启动间隙	融合操作，减少同步
通信与计算重叠度低	阻塞式集合通信	使用异步集合通信，与计算重叠
检测到掉队者	负载不均衡	检查各 rank 的数据分布和批次大小
内存复制开销	数据移动过多	使用原地操作，减少 .cpu()/.cuda() 调用

---

输出格式

报告分析结果时：

1. 先给出摘要 - 用一句话描述主要发现
2. 列出主要问题 - 按耗时列出前 3-5 个内核/操作及其百分比
3. 指出瓶颈类型 - 计算受限、内存受限、通信受限或 CPU 受限
4. 提供可执行的建议 - 给出用户可以进行的具体更改
5. 提供深入分析选项 - 询问用户是否需要 HTA 分析以获取更多详情

输出示例：

## Trace Analysis Summary

**Main Finding**: GPU is idle 40% of the time, primarily due to CPU bottleneck (host wait).

### Top GPU Kernels (by time)
| % Time | Calls | Avg (μs) | Kernel |
|--------|-------|----------|--------|
| 35.2%  | 1000  | 450.3    | ampere_sgemm_128x64 |
| 22.1%  | 1000  | 283.5    | volta_nccl_reduce_scatter |
| 15.8%  | 2000  | 101.2    | elementwise_kernel |

### Bottleneck: CPU-Bound
The GPU is waiting for the CPU 40% of the time. This suggests data loading or preprocessing is slow.

### Recommendations
1. Increase DataLoader `num_workers` (currently appears to be 0 or 1)
2. Enable `pin_memory=True` in DataLoader
3. Move preprocessing to GPU if possible

**Want deeper analysis?** I can use HTA to get detailed idle time breakdown and critical path analysis.

---

规则

1. 缓存前先询问 - 不要自动缓存到 Parquet。务必先询问用户。
2. 始终先进行快速分析 - 无需任何依赖
3. 安装前先询问 - 不要假设已安装 HTA；应提出是否使用
4. 具体明确 - 提供确切的百分比、内核名称和建议
5. 提供后续步骤 - 始终询问用户是否需要深入分析
6. 妥善处理错误 - 如果跟踪格式不符合预期，请说明缺少什么