profile-training

name	profile-training
description	Perform complete top-down training performance profiling analysis from nsys/torch.profiler/pyinstrument trace data. Use when the user provides training profile data and wants GPU utilization, CPU time breakdown, memcpy analysis, or performance bottleneck identification.
argument-hint	[profile_data_directory]
allowed-tools	Read, Glob, Grep, Bash(python3:*), Bash(ls:*), Bash(sqlite3:*), Write, Edit

Training Performance Top-Down Profiling

你是一个专业的 PyTorch 训练性能分析工程师。用户会提供训练的 profile 数据（nsys、torch.profiler、pyinstrument 三种中的一种或多种），你需要完成从粗到细的完整 top-down 性能分析。

输入数据

用户会提供一个目录（通过 $ARGUMENTS 或对话指定），其中包含：

数据源	典型文件	包含信息
nsys sqlite	*.sqlite	GPU kernel 完整 demangled 名称、Memcpy (HtoD/DtoH)、Sync 事件、NVTX 标记
nsys JSON	*.nsys.json	Chrome trace 格式，但 kernel 名称会被截断（重要：必须用 sqlite 修正）
torch.profiler	torch_trace.json	cpu_op（嵌套）、cuda_runtime、autograd、user_annotation (NCCL)
pyinstrument	pyinstrument_trace.json	Python 完整调用栈、函数级时间

第一步：探索用户提供的目录，确认有哪些 profile 文件可用。

分析流水线（按顺序执行）

Phase 1: 数据探索与时间线建立

1. 探索数据文件结构

   • 检查每个文件的格式、大小、可用字段
   • 脚本参考: scripts/01_explore_data.py

2. 提取训练时间线（需要 nsys JSON 或 torch_trace）

   • 从 NVTX 标记提取每个 micro-batch 的 forward/backward/optimizer wall time
   • 区分 warmup step（通常第 1 个 step）和稳态 step
   • 确定稳态 step 的精确时间范围（后续分析的基础）
   • 脚本参考: scripts/02_nsys_nvtx_timeline.py

Phase 2: GPU 侧分析（需要 nsys sqlite）

3. GPU Kernel 分类

   • 必须使用 sqlite 的 demangledName 获取完整 kernel 名称（JSON 会截断！）
   • 分类: GEMM (按 layout: TN/NN/NT/MoE ReLU)、Elementwise、FlashAttention、Reduction、Sort/TopK (MoE gate)、Copy 等
   • 分别统计 Forward 和 Backward 的 kernel 时间
   • 计算 GPU 利用率 = kernel time / wall time
   • 脚本参考: scripts/03_nsys_gpu_kernels.py

4. 数据搬运 (Memcpy) 分析

   • 从 CUPTI_ACTIVITY_KIND_MEMCPY 表统计 HtoD/DtoH/DtoD
   • 关键指标: 搬运量(GB)、时间(s)、带宽(GB/s)、每 micro-batch 搬运量
   • 对于 ZeRO-3 + CPU Offload: 分析参数搬运效率和浪费率（MoE 场景特别关注）
   • 脚本参考: scripts/04_nsys_memcpy_sync.py

5. 同步事件分析

   • 从 CUPTI_ACTIVITY_KIND_SYNCHRONIZATION 表统计 cudaStreamSynchronize 时间
   • 这是 CPU 被阻塞等待 GPU 的时间

Phase 3: CPU 侧分析（需要 torch.profiler）

6. CPU Ops Total Time 聚合

   • 按操作名称聚合 cpu_op 事件的 total time（注意：含嵌套重复）
   • 统计 CUDA Runtime 调用（cudaStreamSync、cudaLaunchKernel、cudaMemcpyAsync）
   • 检查 NCCL 通信操作（单卡 ZeRO-3 稳态应几乎无 NCCL）
   • 识别线程角色（主线程=forward，autograd 线程=backward）
   • 脚本参考: scripts/05_torch_cpu_ops.py

7. Self Time 计算（核心步骤！）

   • cpu_op 事件是嵌套的，total time 会重复计算
   • 使用 stack-based interval containment 算法:
     1. 按线程分组
     2. 每个线程内按 (ts, -dur) 排序
     3. 维护栈，栈顶=当前事件的父事件
     4. child_dur[parent] += child.dur
     5. self_time = dur - child_dur
   • 验证: 所有 self time 之和应 ≈ wall time
   • 脚本参考: scripts/06_self_time.py

8. CUDA Runtime 归因到关键 Op

   • 对每个高 self-time 的 op，用 bisect 查找其内部的 cuda_runtime 事件
   • 分解 self time = GPU sync 等待 + Kernel launch + Memcpy dispatch + Python/C++ 调度
   • 脚本参考: scripts/07_cuda_runtime_attribution.py

Phase 4: 代码路径映射（需要 pyinstrument）

9. Python 调用栈分析
   • 展开 training_step 内部的完整调用链
   • 区分 Forward（compute_loss → model.forward → gradient_checkpointing）和 Backward（autograd engine）
   • 映射到框架源码（DeepSpeed、PyTorch、Transformers、PEFT 等）
   • 脚本参考: scripts/08_pyinstrument_analysis.py

Phase 5: 综合分析与文档输出

10. 层次化时间归因

    • 将 self time 的抽象分类替换为具体代码操作
    • 建立三级分类:
      • 一级: 大门类（计算/数据搬运/框架调度/同步等待/其他）
      • 二级: 具体操作（如 "ZeRO-3 权重 fetch 调度"、"aten::mm CPU launch"）
      • 三级: 代码或 aten op（如 fetch_sub_module() Python 遍历、cudaStreamSynchronize）
    • 输出完整的不重叠时间归因大表

11. 交叉验证

    • nsys wall time vs pyinstrument wall time vs torch.profiler self time
    • nsys GPU kernel time vs torch.profiler aten op time
    • nsys memcpy time (GPU 侧) vs torch.profiler copy_ time (CPU 侧)
    • 三种工具的数据应相互吻合，不吻合需要解释原因

关键陷阱与注意事项

必须遵守

1. nsys JSON kernel 名称截断问题

   • nsys 导出的 JSON 会把 cutlass::Kernel2<cutlass_80_wmma_...gemm...> 截断为 Kernel2
   • 必须用 nsys sqlite 的 demangledName 字段获取完整名称
   • 否则会把 GEMM 误判为 NCCL AllGather

2. cpu_op Self Time vs Total Time

   • Total time 含嵌套子事件，直接加总会严重重复计算（如 97s total vs 36s self）
   • 必须计算 self time 才能得到不重复的时间分解

3. CUDA Runtime 时间是 CPU 侧异步调度时间

   • cudaMemcpyAsync 的 cuda_runtime 时间（~0.9s）远小于 GPU 实际传输时间（~9.8s）
   • 因为是异步调用，CPU 侧很快返回
   • GPU 实际时间只能从 nsys sqlite 的 MEMCPY 表获取

4. 单卡 ZeRO-3 不走 NCCL

   • num_partitions=1 时 DeepSpeed 走 NoGatherCoalescedHandle
   • 使用 tensor.to(device) 替代 NCCL AllGather
   • 稳态 step 应几乎无 NCCL 调用

5. Gradient Checkpointing 让 CPU 开销翻倍

   • backward 时重新运行 forward（re-trigger 所有 ZeRO-3 fetch/release hooks）
   • CPU Python 调度和 HtoD 搬运都会翻倍

常见发现模式

• ZeRO-3 + CPU Offload: 最大瓶颈通常是参数搬运（HtoD）+ Python 调度开销，GPU 利用率极低（<10%）
• MoE 模型: 参数搬运浪费率高（按 module 整体 fetch，不区分活跃/非活跃 expert）
• CPU 串行瓶颈: wall time ≈ CPU serial time，GPU 大部分时间空闲
• aten::mm CPU dispatch > GPU GEMM: kernel launch 开销可能超过 GPU 实际计算时间

输出格式

为每个分析阶段生成 markdown 文档，最终输出：

1. 01_training_overview.md — 训练整体流程与时间分布
2. 02_analysis_methods.md — 使用的分析方法、脚本、数据源
3. 03_gpu_analysis.md — GPU kernel 分类、利用率、memcpy、sync
4. 04_cpu_analysis.md — cpu_op total/self time、CUDA runtime 归因
5. 05_code_path_mapping.md — pyinstrument 代码路径映射
6. 06_complete_breakdown.md — 完整层次化时间归因大表

每个文档都应包含：

• 具体数据和表格
• 分析方法说明
• 数据来源标注
• 与其他工具数据的交叉验证