PyTorch 性能分析入门指南
Hugging Face Blog··作者 Aritra Roy Gosthipaty
关键信息
文章使用了一个最小脚本 01_matmul_add.py,并在 NVIDIA A100-SXM4-80GB GPU 上运行,展示了如何用 torch.profiler.record_function 和 torch.profiler.profile 包装代码。文章还说明,用户需要同时查看 profiler 表格和 Chrome trace,包括 CPU 轨道、GPU 轨道以及它们之间的空隙,才能理解在应用 torch.compile 后哪些部分发生了变化、哪些没有变化。
资讯摘要
Hugging Face 这篇文章开启了一个关于 PyTorch 性能分析的系列,而且明确面向初学者。文章的核心观点是:如果不能进行性能分析,就无法进行优化,因此 profiling 对提升每秒 token 数、降低推理延迟、以及解释训练循环为什么比预期更慢都非常重要。作者也承认,profiler trace 往往很吓人,因为里面充满了密集的彩色块和陌生的事件名称,所以这个系列的目标就是降低阅读门槛。第一部分选择了最简单的例子:矩阵乘法后接一个偏置加法。示例脚本名为 01_matmul_add.py,并在 NVIDIA A100-SXM4-80GB GPU 上运行,文章还建议在正式采集前多运行几次以完成 GPU 预热。
文中展示了如何使用 torch.profiler.record_function 给代码打标签,以及如何用开启 CPU 和 CUDA 活动的 torch.profiler.profile 上下文管理器包住要分析的代码。随后,读者会被引导去查看 profiler 表格和 trace 视图,并注意 CPU 轨道、GPU 轨道以及它们之间的空隙。文章还先解释了两个基础概念:GPU kernel 是在 GPU 上并行运行的程序,而 CPU 负责调度和启动这些 kernel。系列后续内容会逐步扩展到 nn.Linear 和小型 MLP,最后再延伸到使用 transformers 的大型语言模型,并借助 trace 来解释优化以及 torch.compile 带来的变化。
资讯正文
你无法分析的内容,也就无法优化。
无论你是想从大型语言模型(LLM)中榨出更多每秒 token 数,还是想把推理延迟再削掉几毫秒,抑或只是想弄清楚为什么你的训练循环跑得比规格说明书承诺的更慢,最终都绕不开 profiling(性能剖析)。
问题在于,profiling 的入门门槛很高。那些 trace 看起来像一堵堵密密麻麻的彩色矩形墙。事件名称也常常让人望而生畏。大多数教程都默认你已经会读它们。所以即便我们知道自己应该做 profiling,打开一份 trace 也常常像一件最好留到以后、或者留给别人去做的苦差事。本文,以及由它开启的系列文章,就是我们试图降低这道门槛的尝试。
这是《Profiling in PyTorch》系列的开篇文章。这个系列会循序渐进地建立阅读 profiler trace 的能力,并用它来驱动优化。计划如下:
- 第 1 部分(本文):从最简单的操作开始,一个矩阵乘法再加上一次偏置项相加,学习如何读懂 profiler 返回给我们的内容。
- 第 2 部分:扩展到 nn.Linear 和一个小型 MLP,用 trace 来解释为什么要做优化,并顺带看看底层的 kernels。
- 第 3 部分:把前面的内容综合起来,应用到使用 transformers 的大型语言模型上。
我们会以初学者的视角记录这段旅程。除了一些基础的 PyTorch 知识之外,不需要任何前置条件。把这篇文章当作一篇轻松的阅读材料,其中会穿插一些“Aha!”时刻。本文的结构特意采用提问式:我们打开一份 trace,问一句“等等,为什么会这样?”,然后一路追问答案,直到某些东西突然对上了。读完之后,你应该会知道:
- 如何设置 torch.profiler,以及它到底会返回什么;
- 如何阅读 profiler 表格和 trace(CPU 轨道、GPU 轨道,以及它们之间那些可疑的空隙);
- 从一次 Python 调用一路到底层 CUDA kernel 的事件链条;
- 当你在外面套上一层 torch.compile 时,会发生什么变化,更有意思的是,又有哪些东西并不会变化。
在开始之前,先给出两个定义,它们会让下面的内容更容易理解:
- GPU kernel 是一种程序,它会在 GPU 的许多线程上并行运行。
- CPU 负责调度并启动这些 kernels。
你通常不需要自己编写 GPU kernel;当你使用 PyTorch 的某个操作时,它会被自动翻译成一个或多个在 GPU 上完成任务的 kernels。
有了这两个概念打底,我们就开始提问吧。
下面是本文使用的完整脚本:01_matmul_add.py。我们建议你在另一个标签页中打开这个脚本,并逐步跟着代码走一遍。运行这些脚本时,我们使用的是 NVIDIA A100-SXM4-80GB GPU。
正如 Sara Hooker 博士那句妙语所说,正如我们主要由水构成一样,深度神经网络主要由矩阵乘法构成。既然矩阵乘法如此基础,如果我们不从它开始这段 profiling 之旅,那可就太可惜了。
def fn(x, w, b):
return torch.add(torch.matmul(x, w), b)
矩阵乘法再加上矩阵加法,模拟的是神经元中权重和偏置之间的相互作用。这个“加法”(这里顺便玩个双关)会帮助我们理解它如何为后文的编译铺平道路。
要进行 profiling,我们将使用 torch.profiler 模块。涉及的步骤如下:
准备好要分析性能的代码(这里是 fn,它封装了矩阵乘法和矩阵加法)——为算法添加标注。虽然这完全是可选的,但我们建议这样做。record_function 会把我们的函数标记为 matmul_add,这样在跟踪信息中就更容易导航(我们稍后会提到这一点)
def step():
with torch.profiler.record_function("matmul_add"):
return fn(x, w, b)
- 用 torch.profiler.profile 上下文管理器包裹代码
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()
- 导出性能分析结果
# profiler 表格
prof.key_averages().table(sort_by="cuda_time_total", row_limit=15)
# profiler 跟踪
prof.export_chrome_trace(trace_path)
profiler 会导出两种不同的产物:
- profiler 表格:提供该算法的统计摘要。它回答的是“什么最耗时”。这对于找出热点非常有帮助。热点指的是最耗时的事件、可能成为流水线瓶颈的事件,或者被触发很多次的事件。
- profiler 跟踪:提供按时间展开的执行视图。它回答的是“一个操作在什么时候以及为什么发生”,并描绘 CPU 和 GPU 上发生的活动。当我们想调查被启动的 kernel、启动它们时的任何延迟、CPU 与 GPU 活动之间的重叠等情况时,这会很有帮助。
让我们通过第一次执行来看看这两者的实际效果。(这里是完整的 01_matmul_add.py 脚本)
建议在带有 GPU 的机器上运行这个脚本。
uv run 01_matmul_add.py --size 64
如果你在(有 GPU 的)机器上运行上述脚本,你会发现一个 traces/01_matmul_add 文件夹,其中包含这两个产物:
64_bf16_cold_eager.json
64_bf16_cold_eager.txt
.txt 文件保存的是 profiler 表格。打开该文件后,如图 1 所示,首先会看到一个大表格,其中第一列由在 profile 作用域内被触发的事件组成。
其他列与事件在 CPU 或 GPU 上所花费的时间有关,也可能与 torch.profiler.profile 的 activities 中指定的其他设备有关。看看哪些事件最耗时,并尝试直观地理解该事件是否确实应该花这么多时间。另一个重要的点是查看 “# of Calls” 列,它说明了该事件被触发了多少次。
顺便说一下,我们也来谈谈 “Self CPU/CUDA” 与 “CPU/CUDA total” 之间的区别。“Self” 列只衡量事件自身内部耗费的时间,不包括其子事件。“total” 列则把该事件及其所有子事件的时间一起算上。所以如果你查看 matmul_add 的 “CPU total”,它包含了自身耗时以及它触发的子事件耗时。这是一个需要注意的重要细节。
如果你看表格最后两行,会发现 profiler 告诉我们:
Self CPU time total: 2.314ms
Self CUDA time total: 23.104us
CPU 时间以毫秒为单位
而 GPU 时间单位是 us。
为了便于理解,GPU 上花费的时间(内核 ampere_bf16_s16816gemm...)不到 CPU 上花费时间(matmul_add 操作)的 1%。GPU 大部分时间都处于空闲状态,这显然是一个危险信号。之所以会这样,是因为 GPU 能非常快地计算一个小矩阵乘法,因此我们的代码大部分时间都花在准备内核、将其启动到 GPU 上、发送要相乘的数据以及收集结果上。这个概念被称为“开销受限”算法。
摆脱这种状态的最简单方法是使用更大的矩阵乘法。
uv run 01_matmul_add.py --size 4096
图 2 的最后两行是:
Self CPU time total: 4.908ms
Self CUDA time total: 4.495ms
这两个时间单位都是 ms,这意味着仅仅通过增大矩阵乘法的规模,我们就让更多 GPU 时间实际被利用起来了。如果你查看图 2,还会注意到,现在绝大部分 CUDA 时间都由 GPU 内核(ampere_bf16_s16816gemm_..)占用,而不再是启动它的 CPU 操作(matmul_add)。这说明我们确实已经从开销受限转向了计算受限。
接下来,我们进入对调度链的可视化分析,它位于 .json 工件中。你可以将它们上传到 Perfetto UI 中查看追踪,也可以使用 uvx trace-util traces -b traces 直接生成 Perfetto 链接。
在图 3 中,我们看到了矩阵乘法和加法的 profiler 追踪图。这里条形的宽度表示事件持续时间,垂直嵌套表示调用层级,CPU 轨道表示发生在 CPU 上的事件,而 GPU 轨道则显示实际的内核执行。你还可能会注意到那些空白区域,它们就是等待或空闲时间。
脚本使用的是默认配置,分别是:
- size 64:输入、权重和偏置的尺寸为 (64, 64)
- dtype bf16:数据类型是 bfloat16
- no compile:我们没有对 torch 操作进行编译
- no warmup:在 profiling 之前,我们没有先对 GPU 进行预热
在 Perfetto 中,我们建议使用键盘来更快地浏览追踪图。可以使用 “W A S D” 来导航追踪视图。
图 4 中有两条轨道,一条是 CPU 活动,另一条是 GPU 活动。在 CPU 轨道中,可以看到三个 profile 步骤(从 ProfilerStep#2 开始)。这来自如下调度设置:
schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
wait 会跳过噪声较大的初始化阶段(ProfilerStep#0),warmup 会在不记录的情况下运行 profiler(ProfilerStep#1),而 active 则是最终会显示在追踪图中的部分。脚本中使用的 schedule 可以在这里找到。
让我们戴上侦探帽,检查这条追踪,并提出一些问题。
在图 5 中,我们注意到 ProfileStep#2 比其他步骤花费更多时间,仔细观察后,你也会在 matmul_add 标注中看到类似的模式。关键证据藏在标注内部,而不是标注本身:
图 6 中显示的约 228 µs 是进入 record_function("matmul_add") 到 PyTorch 实际调度 aten::matmul 之间的“空窗期”。
这可能有多种原因,包括工作区分配、cuBLAS(NVIDIA 专有的、用于执行基本线性代数运算的 GPU 加速库)启发式策略,或者是惰性模块加载。我们可以选择忽略它,或者在开始分析之前再运行一些热身步骤(这也是标准做法)。
在性能分析中,热身指的是在真正开始分析之前先把这些事件运行几次。GPU 预先完成的工作(包括上面提到的这些内容)都是一次性的开销,我们并不希望把它们纳入分析。在我们的示例中,有两个热身阶段:一个是在进入 profiler 之前,我们实际上先循环执行函数;另一个是在 profiler 内部完成,这通过 warmup 参数实现。在这一节中,我们已经启用了实际迭代以及调度器。
uv run 01_matmul_add.py --warmup
在图 7 中,我们看到每个 profile step 的耗时都很接近,但这并不意味着我们已经成功优化了这些一次性开销。我们之所以先进行热身,是为了让这些开销不被纳入分析。我们认为,如果在这里突然结束这一节而不给出解决思路,对读者是不公平的,所以这里提供一个链接,介绍如何进一步优化启动开销。
在图 8 中,我们看到 CPU 和 GPU 轨道之间大约有 2.5 ms 的偏移:这是 CPU 提交 CUDA 内核之后,到它们真正开始执行之间的延迟。人们可能会认为,热身阶段加上调度器的 wait
和 warmup
应该能让 GPU 一直保持忙碌,并减小这种偏移。
为了弄清楚到底发生了什么,我们稍微修改一下调度器:
- schedule = torch.profiler.schedule(wait=1, warmup=1, active=3, repeat=1)
+ schedule = torch.profiler.schedule(wait=0, warmup=0, active=3, repeat=1)
图 9 显示,在 GPU 轨道上,在任何操作之前就出现了一个 Activity Buffer Request。
让我们再放大一点看看。
放大 GPU trace 之后,我们注意到,ProfileStep#0 的 matmul 和 add 内核(其 CPU trace 在图中不可见)是一个接一个执行的,而 ProfileStep#1 的内核之间则留出了一段空白窗口。对此最合理的解释是,缓冲区发生了溢出,并且在内核执行期间又发出了另一个 buffer request(即请求在 GPU VRAM 中分配一些内存)。
排除其他可能性的最好办法,是分析更多迭代,看看 trace 的其他部分是否也会出现类似的空白窗口。为此,我们使用 active=20 来运行。
如图 11 所示,我们在 ProfileStep#1 中看到类似的趋势。这与我们之前的发现一致,因此我们可以放心地得出结论:这确实是另一个 buffer request。
在图 12 中,我们看到了嵌套的 CPU 调用。这是一个重要的可视化方式,它能让人理解一条 dispatch 链到底是什么样子。
我们从 ProfileStep#<id> 开始,它封装了 profiling step。由于我们对该步骤做了标注,因此可以看到 matmul_add
这一行。matmul_add
由两个 aten
调用组成,一个用于矩阵乘法,一个用于矩阵加法。
aten::matmul
是用户层面的 PyTorch matmul 调用最终落到的 ATen 级别 dispatch。aten::mm
这是 2D 矩阵乘矩阵的乘法后端。
非常有意思的是,如果我们给矩阵加上 batch 轴,PyTorch 会如何调用 aten::bmm(批量矩阵乘法)。让我们稍微绕个弯,看看 aten::bmm 的实际运行情况。
- x = torch.randn(args.size, args.size, device=device, dtype=dtype)
- w = torch.randn( args.size, args.size, device=device, dtype=dtype)
- b = torch.randn(args.size, args.size, device=device, dtype=dtype)
+ # adding a batch size of 8
+ x = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
+ w = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
+ b = torch.randn(8, args.size, args.size, device=device, dtype=dtype)
在图 13 中,当给输入添加 batch 轴后,aten::matmul 现在会封装一系列其他必要的 CUDA 运行时调用,以及 aten::bmm(而不是 aten::mm)。这也提示了 cuBLAS 为了为程序调度正确的、最合适的 kernel 需要做出的启发式判断。
在本文剩余部分中,除非另有说明,我们都将使用简单的 2D 矩阵。
我们注意到,对于 aten::mm,有两个 CUDA Runtime 调用,分别是 cudaOccupancyMaxActiveBlocksPerMultiprocessor(在图 14 中用框标出)和 cudaLaunchKernel;而对于 aten::add,则只有 cudaLaunchKernel。cudaOccupancyMaxActiveBlocksPerMultiprocessor 是一个规划调用,完全发生在 CPU 端。它会询问:“给定一个 kernel 函数、选定的 block 大小,以及选定的动态共享内存大小,这个 kernel 有多少个 block 可以同时驻留在一个 SM(Streaming Multiprocessor)上?”
这就引出了一个问题:为什么矩阵乘法需要规划,而 add 不需要?
要理解这一点,我们必须看看 kernel 的资源占用情况。如果你点击 GPU kernels,就可以查看各自 kernel 的资源占用。
在图 15 中,我们注意到,对于矩阵乘法来说,每个线程的寄存器数量和共享内存是动态的(取决于矩阵大小)。cuBLAS 提供了数百种 kernel 变体,每一种都有一条由启发式驱动的启动路径,而这条路径需要运行时获得硬件容量信息。occupancy 查询就是这种启发式的一部分。从概念上讲,我们可以把 GPU 加速的矩阵乘法理解为在处理彼此独立的 tile:我们使用多少个 tile、每个 tile 需要多大,取决于矩阵和硬件。现代算法要复杂得多,但这仍然是一个很好的参考框架。
从图 16 可以看到,加法的资源占用显示为 32 个寄存器和 0 共享内存。这完全是轻松可容纳的。这里没有什么需要查询的,因为没有任何硬件资源会限制 occupancy。这个 kernel 在设计上就是资源消耗很低的。
你可以把这一点作为阅读任何 trace 时的快速诊断方法。扫描 CPU 轨道中的 cudaOccupancyMaxActiveBlocksPerMultiprocessor。每出现一次,就说明有一个“重量级、可自适应调度”的 kernel,通常是 GEMM(GEneral Matrix Multiplication)、conv,或类似算子。没有前置 occupancy 查询的 kernel,则属于逐元素/归约这一类,PyTorch 会机械地启动它们。
cudaDeviceSynchronize
CPU 会一直阻塞,直到该设备上的所有 GPU 工作完成。分析器会在活动窗口结束时发出这个同步,以便刷新事件。如果没有它,内核耗时就会缺失。
一个覆盖 26 µs 实际 GPU 工作的 1.78 ms 同步,说明这次运行有 98% 的时间都处于空闲状态。这就是教科书式的开销受限症状。
我们已经从上面的分析器表格分析中知道,把更大的矩阵传给算法,会使它从开销受限区域转向计算受限区域。
让我们运行命令并更深入地查看这些 traces。
uv run 01_matmul_add.py --size 4096 --warmup
在图 17 中,我们注意到 ProfileStep#3 的 matmul 内核
在 GPU 上比其他步骤耗时更长。值得特别注意的是,启动的其他内核完全相同,这意味着其中没有涉及 cuBLAS heuristics。不存在调度间隙,CPU 发起调用也很正常,这也不是分析器伪影。
图 17 中的这个 trace 提供了一个很有用的观点,而这在理想化示例中很容易被忽略:即使在相同硬件环境、相同数据上运行相同代码,内核运行时间也不是常数。
让我们通过对脚本做一点小改动,把这一点具体化。我们把迭代运行 20 次,捕获每个步骤。
+ schedule = torch.profiler.schedule(wait=0, warmup=0, active=20, repeat=1)
- for _ in range(5):
+ for _ in range(20):
图 18 显示了类似的发现。虽然每个内核都完全相同,但它们的耗时却不同。这种不同的计算时间可以归因于很多原因:
- GPU 在空闲与加速状态下的时钟频率
- GPU 发热
- GPU 电源管理
- 驱动端的维护工作
只看平均值的读者会得出结论:一次 matmul 约耗时 ~1 ms(5 次的均值 = 1084 µs);而看了 trace 的读者会发现,matmul 通常耗时约 ~580 µs,除非 GPU 发脾气。那是两种截然不同的心智模型,而且只有一种是正确的。
使用 torch.compile
一直让我们惊叹。人们写下普通的 eager PyTorch 代码,而 PyTorch 会尝试捕获张量密集区域,把它们转成图,进行优化,然后运行生成的代码。默认后端通常是 TorchInductor
,而整体流程大致如下:
TorchDynamo
将 Python 执行捕获为 FX 图AOTAutograd
在涉及梯度时准备前向/反向图Inductor
把图降级为优化后的 CPU 或 GPU 代码。
在这一部分,我们讨论编译并查看分析器 trace。
uv run 01_matmul_add.py --size 4096 --warmup --compile
args.compile
标志会触发下面这段代码:
fn = torch.compile(fn) if args.compile else fn
在图 19 中,我们看到了名为 Torch-Compiled Region: 0/0
的新 CPU 行,这指向了正在使用已编译函数。
查看图 20 时,我们会问一个问题:我们是否真的把乘法和加法操作融合成了一个操作?
这就是图级别的算子融合。Inductor 把我们的 torch.add(torch.matmul(x, w), b)
重写成了一个单独的 aten::addmm(b, x, w)
调用。这里需要注意的重要一点是,它并没有生成一个新的融合 CUDA kernel。实际的 GPU 工作仍然是 ampere_bf16_s16816gemm_bf16_128x256_ldg8_f2f_stages_64x3_nn
,也就是 eager 模式所使用的同一个 cuBLAS kernel。所以这里的“融合”发生在分发器层面,而不是在 kernel 层面。
PyTorch 提供了 torch.addmm 函数,它做的正是我们分两步完成的操作,也就是先乘再加。我们鼓励读者查看这个函数的 trace,并在下面的评论里写下你的观察!
虽然我们在理论上知道编译函数时会发生什么,但亲眼看到它实际运行同样重要。让我们看看 CPU 侧的层次结构,它反映了 torch.compile 的运行时架构。
TorchDynamo Cache Lookup 是 Dynamo 检查当前调用是否仍然与编译时一致的地方,这些条件包括输入形状、dtype、设备以及 tensor 元数据。如果有任何不匹配,Dynamo 就会重新编译。即使在编译之后,这个开销也会在每次调用时支付。
Torch-Compiled Region 是“进入”编译版本的包装器。AOTDispatcher Runtime Wrapper Prologue 是 AOT Autograd 的运行时包装器。即使这里我们不需要梯度,AOTDispatcher 也总是在调用栈中,负责处理 tensor 元数据、view 跟踪,并且在 requires_grad 为 true 时会为反向传播做准备。
## Call CompiledFxGraph 是实际生成代码运行的地方。"CompiledFxGraph" 后面的字符串是 FX graph 的内容哈希。它在这三个活跃步骤中都相同,确认了缓存命中。
你可以在磁盘上 /tmp/torchinductor_<user>/fxgraph 下找到生成的代码,并按这个哈希进行键控;当你想阅读 Inductor 实际生成的 Triton/C++ 代码时,这很有用。
查看图 21 中的 trace 时,我们非常高兴地注意到每一步只有一个 cudaLaunchKernel。这一观察与我们在 GPU trace 中看到的情况直接矛盾。实际上每一步仍然启动了两个 kernel,分别是 Memcpy DtoD(Device -> Device)和 GEMM。回到 CPU trace 后,我们注意到自己完全漏看了 cudaMemcpyAsync 的分发。
addmm 计算 out = α·A·B + β·C,而 cuBLAS 的带 bias add epilogue 的 GEMM 会写入一个目标缓冲区,而这个缓冲区里需要事先包含 bias。epilogue 可以理解为 GEMM 之后发生的所有操作。在深度学习领域,我们不断提出各种 GEMM-Epilogue,比如激活函数、bias addition、归一化等等。这也是为什么会有 cuBLAS 的 GEMM-with- kernel 变体。
如果你对 torch.compile 使用不同的 mode,你会注意到启动了不同的 kernel 变体。你可以亲自试试看,并在下面留下你观察到的评论!
因此,Inductor 生成的代码做的是:
out = copy(C)
← 这就是 DtoD memcpy(32 MB,耗时约 33 µs)out = α·(A·B) + β·out
← 带 α=β=1 的 GEMM
,将 bias add 融合进写回阶段。
其结果在数学上仍然是相同的。bias add 并不是免费的,因为我们先要支付一次 memcpy 的开销,然后还要承担一个稍微更昂贵的 GEMM epilogue。
人们可能原本希望看到的那种融合,也就是 x·w + b
(这里 out = α·A·B + β·C
把它看作一个单独的 kernel、没有额外的内存流量,并不是实际发生的情况。Inductor 保留了这两个会触及内存的操作,只是把 bias 复制重命名为 memcpy,把加法重命名为 GEMM epilogue。
真正融合的实现会跳过 memcpy。这正是 FlashAttention 风格的手写 kernel 所做的事情,也是 Inductor 通过 Triton 代码生成可以做到的;但对于一个 4096×4096 的 bf16 matmul,Inductor 显然认为“使用 cuBLAS,通过 epilogue 进行 bias 处理”是最好的路径。
在比较 eager 运行和编译后运行时,最容易忽略的就是这一点:
编译在 CPU 上每一步大约要贵 2 倍。这是因为每次调用都会完整走一遍 Dynamo > AOTAutograd > Inductor 这条栈,而我们本来就已经有同样的 aten::addmm 分发开销了。这个编译流水线是为有几十个操作的机器学习模型设计的,在那种场景下每次调用的额外开销可以被摊薄(但对于单个操作来说,这就是一种额外负担)。
它有一个 amode 参数。读者可以把这当作一个课后作业:去阅读文档,并想出一个能够降低 CPU 开销的 amode。🤗
下面是我们讲到的模式的快速参考。思路是:如果你在 trace 里看到这个,那么它通常意味着什么。
我们从一个很小的 matmul + add 开始,把它当作学习如何阅读 PyTorch profiler 的借口。在这个过程中,我们掌握了一些很适用于更大规模工作负载的思维模型。这是 Profiling PyTorch 系列的第一站。在后续文章里,我们会逐步离开这个两个算子的玩具示例,沿着复杂度阶梯往上走,观察更大的构建块,并最终看向真实模型。
感谢 Noe Flandre、Suvaditya Mukherjee 和 Vidit Ostwal 对这篇文章早期草稿所做的审阅!
来源与参考
收录于 2026-05-30