PyTorch 性能分析：torch.profiler 初学者指南

PyTorch 性能分析：torch.profiler 初学者指南

了解 torch.profiler 进行性能优化

性能分析是优化大型语言模型（LLMs）和深度学习流水线的关键第一步，因为它可以帮助开发者判断模型是受开销限制还是计算限制。torch.profiler 模块提供了两种主要的诊断产物：profiler table，提供了“什么”占用最多时间的统计摘要；以及 profiler trace，提供了在 CPU 和 GPU 通道上操作“何时”以及“为何”发生的时间视图。

识别瓶颈：开销受限 vs. 计算受限

性能瓶颈通常通过比较 profiler table 中 CPU 与 GPU 的耗时来揭示。

  • Overhead-Bound: 当 Self CPU time total 明显高于 Self CUDA time total（例如 CPU 耗时为毫秒级而 GPU 为微秒级）时，算法属于开销受限。这表明 CPU 在准备和启动 kernel 上花费的时间多于 GPU 的执行时间。小规模矩阵乘法常出现此情况。
  • Compute-Bound: 当 Self CPU time total 与 Self CUDA time total 均在毫秒级且相当时，算法属于计算受限。这是高性能计算的理想状态，GPU 成为主要瓶颈。

增大工作负载规模（例如从 64×64 矩阵提升到 4096×4096 矩阵）通常会将模型从开销受限的 regime 转变为计算受限的 regime。

分析 CPU 与 GPU 调度链

PyTorch 操作遵循从 Python 调用到底层 CUDA kernel 的特定调度链。典型的矩阵乘法与加法序列如下：

ProfileStep → record_function (user annotation) → aten::matmul (ATen-level dispatch) → aten::mm (2D matrix-matrix multiply backend) → cudaLaunchKernel.

CPU 通道关键指示

  • Cold-Start Overhead: 第一次 ProfileStep 通常比后续步骤更宽，因为涉及工作空间分配、cuBLAS 启发式选择以及惰性模块加载。可通过热身迭代来缓解。
  • CUDA Occupancy Queries: 在 cudaLaunchKernel 之前出现 cudaOccupancyMaxActiveBlocksPerMultiprocessor 表明这是一个“重量级” kernel（如 GEMM 或卷积）。CPU 正在查询驱动以根据硬件容量确定最佳块大小。
  • Resource-Light Kernels: 元素级或归约 kernel 通常没有占用查询，因为它们的资源占用固定且很小。
  • Synchronization: 跟踪末尾出现的长时间 cudaDeviceSynchronize 往往是 profiler 刷新事件的过程；其时长反映 CPU 等待 GPU 完成挂起工作所耗费的时间。

GPU 通道关键指示

  • Activity Buffer Requests: kernel 之间的空隙或 CPU 与 GPU 通道之间的初始偏移常由 profiler 分配或重新填充自身事件缓冲区导致。
  • Runtime Variance: 相同的 kernel 在不同步骤可能表现出不同的运行时，原因包括 GPU 时钟波动、热量、功耗管理或驱动维护等因素。

torch.compile 的影响

使用 torch.compile 可通过 TorchDynamo、AOTAutograd 与 TorchInductor 将即时 PyTorch 代码转化为优化图。对已编译函数进行分析会揭示若干架构变化。

调度器级别融合

对于 torch.add(torch.matmul(x, w), b) 之类的操作，torch.compile 在图层面进行算子融合，将独立的 aten::add 与 aten::mm 调用替换为单一的 aten::addmm 调用。但这仅是 调度器级别 的融合，而非 kernel 级别。GPU 仍会执行一次 Memcpy DtoD（用于将偏置写入目标缓冲区），随后运行带有偏置加法尾部的 GEMM kernel。

运行时架构与开销

已编译区域在 CPU 端引入了特定的层级结构：

  1. TorchDynamo Cache Lookup: 验证输入的形状、数据类型和设备是否匹配缓存的编译结果。此步骤在每次调用时都会执行。
  2. Torch-Compiled Region: 进入已编译版本的包装器。
  3. AOTDispatcher Runtime Wrapper Prologue: 处理张量元数据和视图追踪。
  4. Call CompiledFxGraph: 执行生成的代码（通过内容哈希标识）。

对于极小的操作，torch.compile 可能会增加 CPU 开销，因为遍历 Dynamo → AOTAutograd → Inductor 堆栈的成本超过了融合带来的收益。随着模型中操作数量的增加，这部分开销会被摊薄。

跟踪读取摘要表

观察 可能含义
Self CPU \gg Self CUDA 开销受限；增大批量大小或融合操作
cudaOccupancy... before launch 重量级、适配性启动的内核（GEMM/卷积）
aten::matmul → aten::bmm 对 3D 以上张量的批量矩阵乘法
Torch-Compiled Region 执行位于 torch.compile 优化块内
Memcpy DtoD before GEMM 编译模式下 addmm 尾部的偏置拷贝

Sources