Profiling in PyTorch (Part 1): torch.profiler 初学者指南深度解读

背景

在深度学习领域，有一句至理名言：“无法剖析（Profile）的东西，就无法优化。”

无论你是试图从大语言模型（LLM）中压榨出更多的每秒令牌数（tokens per second），还是希望在推理过程中削减几毫秒的延迟，亦或是仅仅想弄清楚为什么你的训练循环速度远低于规格说明书上的承诺，最终的路径都指向了性能剖析（Profiling）。

然而，性能剖析的学习曲线相当陡峭。生成的追踪记录（Traces）通常是密集且令人困惑的彩色矩形墙，事件名称也往往带有威慑力。大多数教程都默认读者已经能够阅读这些复杂的图表，导致即使我们知道应该进行剖析，打开追踪文件也往往被视为一种负担，被推迟到“以后再说”或交给其他人处理。

本文旨在降低这一门槛。作为 Hugging Face Blog 发布的 "Profiling in PyTorch" 系列的第一部分，我们将以初学者的视角，通过问答式结构，逐步构建阅读 Profiler 追踪记录的技能，并利用这些洞察驱动优化。

核心内容

1. 基础概念与前置定义

在深入代码之前，需要明确两个关键概念，这将贯穿全文：

• GPU Kernel（GPU 内核）：在 GPU 的许多线程上并行运行的程序。
• CPU 调度：CPU 负责调度并启动这些内核。

通常情况下，开发者无需手动编写 GPU 内核。当使用 PyTorch 操作时，它会自动被翻译为一个或多个在 GPU 上执行任务的内核。

2. 剖析对象：矩阵乘法与加法

为了模拟神经网络中权重和偏置的交互，我们选择一个基础操作作为剖析对象：矩阵乘法（Matmul）加上矩阵加法（Add）。

def fn(x, w, b):
    return torch.add(torch.matmul(x, w), b)

这个看似简单的操作将帮助我们理解后续编译（Compilation）带来的变化。

3. 使用 torch.profiler 进行剖析

PyTorch 提供了 torch.profiler 模块来收集性能数据。完整的剖析流程包括以下步骤：

步骤一：准备代码并添加注解

虽然注解是可选的，但强烈建议添加。使用 torch.profiler.record_function 可以将函数标记为特定名称（如 matmul_add），便于在追踪记录中导航。

def step():
    with torch.profiler.record_function("matmul_add"):
        return fn(x, w, b)

步骤二：包裹剖析上下文

使用 torch.profiler.profile 上下文管理器包裹代码，并指定要监控的活动（CPU 和 CUDA/GPU）。

with torch.profiler.profile(
    activities=[
        torch.profiler.ProfilerActivity.CPU, # CPU 活动
        torch.profiler.ProfilerActivity.CUDA, # GPU 活动
    ],
) as prof:
    # 建议运行多次以预热 GPU
    for _ in range(5):
        step()
    prof.step()

步骤三：导出剖析结果

剖析器会生成两种截然不同的工件（Artifacts）：

1. Profiler Table（剖析表）：

   • 格式：通常为 .txt 文件。
   • 作用：提供算法的统计摘要，回答“什么最耗时”。
   • 用途：识别热点（Hotspots）。热点是指耗时最长、可能是管道瓶颈或触发频率极高的事件。
   • 关键字段解读：
     • Self CPU/CUDA vs CPU/CUDA total：
       • Self（自身）：仅测量事件本身花费的时间，不包括其子事件。
       • Total（总计）：包括事件本身及其所有子事件的时间总和。
     • # of Calls：事件触发的次数。

2. Profiler Trace（追踪记录）：

   • 格式：通常为 .json 文件（Chrome Trace 格式）。
   • 作用：提供时间执行视图，回答“操作何时发生以及为何发生”。
   • 用途：可视化 CPU 和 GPU 上的活动，调查内核启动延迟、CPU/GPU 活动重叠情况等。可通过 Perfetto UI 或 trace-util 工具查看。

4. 案例分析：从 Overhead-bound 到 Compute-bound

为了直观展示剖析结果，我们使用 01_matmul_add.py 脚本，并在 NVIDIA A100-SXM4-80GB GPU 上运行。

场景 A：小矩阵乘法（Size 64x64）

运行命令：uv run 01_matmul_add.py --size 64

生成的追踪记录显示：

• Self CPU time total: 2.314 ms
• Self CUDA time total: 23.104 us

解读： GPU 时间（内核 ampere_bf16_s16816gemm... 执行时间）不到 CPU 时间的 1%。这意味着 GPU 大部分时间处于空闲状态。 原因：这是一个典型的 Overhead-bound（开销主导） 算法。由于矩阵很小，GPU 计算极快，代码大部分时间都花在准备内核、启动内核、发送数据和收集结果上，而非实际计算。

场景 B：大矩阵乘法（Size 4096x4096）

运行命令：uv run 01_matmul_add.py --size 4096

生成的追踪记录显示：

• Self CPU time total: 4.908 ms
• Self CUDA time total: 4.495 ms

解读： 随着矩阵尺寸增大，GPU 时间显著增加，且与 CPU 时间处于同一数量级。此时，大部分 CUDA 时间由 GPU 内核（ampere_bf16_s16816gemm...）占据，而非启动它的 CPU 操作。 结论：我们成功将算法从 Overhead-bound 转变为 Compute-bound（计算主导）。此时，瓶颈真正转移到了计算能力上，而非系统开销。

关键要点

• 剖析是优化的前提：只有通过剖析了解性能瓶颈，才能进行有效的优化。
• 两种剖析工件各有侧重：
  • Table（表）：用于宏观统计，快速定位“最耗时”的操作（热点）。
  • Trace（追踪）：用于微观时序分析，查看 CPU 与 GPU 的协作、延迟及重叠情况。
• 区分 Self 与 Total 时间：
  • Self 时间反映操作本身的效率。
  • Total 时间反映操作及其子操作的总开销。
• 识别 Overhead-bound 现象：当 GPU 时间远小于 CPU 时间（例如小于 1%）时，表明系统开销（启动、数据传输）主导了性能，而非计算本身。
• 规模影响性能特征：小矩阵运算容易陷入 Overhead-bound，增大矩阵规模可将其转化为 Compute-bound，从而更真实地反映 GPU 的计算能力。
• 预热的重要性：在剖析前运行多次循环以预热 GPU，可避免冷启动带来的数据偏差。

意义与影响

这篇文章为 PyTorch 用户提供了一个清晰、低门槛的性能剖析入门路径。它打破了“剖析即复杂”的迷思，通过具体的代码示例和对比实验，让开发者能够直观地理解 CPU 与 GPU 之间的交互机制。

对于深度学习工程师而言，掌握 torch.profiler 的基本用法是迈向高性能模型优化的第一步。理解何时系统是受限于开销（Overhead），何时受限于计算（Compute），是决定后续优化策略（如使用 torch.compile、算子融合或调整数据布局）的关键依据。

本系列后续文章将进一步深入，探讨如何将 nn.Linear 融合为 Fused MLP，以及在大型语言模型（LLM）中如何综合运用这些剖析技巧，最终实现从基础操作到大规模模型的全链路性能优化。