Hugging Face 6 月 11 日发布 PyTorch Profiling 系列第二篇,测试环境里用到了 NVIDIA A100-SXM4-80GB。文章看的不是新模型,也不是新框架,而是一个更容易误读的问题:torch.compile 到底改了什么。
最反常的地方在这里。CPU trace 里一串 aten::t、transpose、view、reshape,看起来很忙;但 GPU 上真正跑的 kernel,可能几乎没变。
这对做 PyTorch 推理、训练优化的人很要紧。因为很多性能报告会把 CPU trace 里的 op 数量,直接当成 GPU kernel 数量。这个错一旦带进优化方案,后面很容易走偏。
单个 Linear:bias 已经折进 GEMM,compile 空间很小
nn.Linear 本质上就是:x @ w.T + b。
但 PyTorch 不会真的先在 GPU 上复制一份转置后的权重,再单独做一次 bias add。Hugging Face 的 profiler trace 指向的是 aten::addmm。bias 会被放进 GEMM 的 epilogue,在结果写回显存前完成加法。
所以,单个带 bias 的 Linear,本来就已经比较紧。
CPU trace 里出现的 aten::t、aten::transpose、view、reshape、as_strided,大多是元数据操作。它们改的是 shape 和 stride,不等于 GPU 上真的搬了一遍矩阵。
这也是很多人读 profiler 时最容易踩的坑:CPU 侧看到多个 aten op,不等于 GPU 侧 launch 了多个 kernel。
放到 torch.compile 上,结论就很清楚。对单个 nn.Linear,compile 前后 GPU 上的 GEMM kernel 基本相同。它能省的主要是部分 CPU dispatch、view 链路,不是把矩阵乘法本身变成了另一种神仙算法。
对工程团队来说,这会影响动作选择。
如果你只是在一个大 Linear 上套 torch.compile,不要急着把它写进“核心加速手段”。更现实的做法是先确认:耗时到底在 GEMM,还是在 Python 调度和小算子链。
GeGLU MLP:真正能省的是 launch 和 pointwise
GeGLU MLP 的情况不一样。它不是一个孤立 Linear,而是多个 Linear 加上激活和乘法。
在 eager 模式下,原文里的预期路径是:3 个线性层对应 3 个 GEMM,GeLU 是 1 个 pointwise kernel,mul 也是 1 个 pointwise kernel。合起来是 5 个 GPU kernel。
这时 compile 才有更明确的空间。
它不需要改写 3 个 GEMM 的本质。它更可能做的是减少 CPU dispatch,并把 GeLU、mul 这类 pointwise 操作合并成更少的 GPU kernel。
| 场景 | CPU trace 里容易看到什么 | GPU 侧更该看什么 | 工程判断 |
|---|---|---|---|
单个 nn.Linear | aten::t、aten::addmm、view 类操作 | 仍是 GEMM-with-bias | compile 收益有限,主要省调度链 |
| GeGLU MLP eager | 3 个 Linear、GeLU、mul | 3 个 GEMM + 2 个 pointwise kernel | 预期约 5 个 GPU kernel |
| GeGLU MLP compile | 3 个 aten::mm 加 fused pointwise | GEMM 名称大体相近,pointwise 可能合并 | 收益来自少 launch、少 dispatch |
这里的主线不是“compile 有没有用”。
它当然有用,但用处有边界。单个大 GEMM 已经由 cuBLAS、CUTLASS 这类库打磨多年,能挤出的空间有限。多个小算子夹在 GEMM 中间时,融合才更容易产生体感收益。
这也解释了为什么同样是 torch.compile,有人测得明显变快,有人几乎没感觉。shape、dtype、算子组合、后端 kernel 选择,都会改变结果。没有这些条件,单说快或慢,都太粗。
性能报告该怎么改:少数 op,不如看 kernel 名称
读这类 profiler trace,我更在意 kernel 名称,而不是 CPU op 数量。
比如 CUTLASS kernel 名称里常能看到 bf16、tile 形状、tn 等信息。tn 说明 GEMM 读取矩阵布局的方式,kernel 可以按转置视图去读权重,并不需要先生成一个真实转置矩阵。
如果 eager 和 compile 下 kernel 名称逐字一致,至少说明 GPU 核心计算路径没有明显变化。此时把加速归因到“GEMM 被优化了”,就站不住。
更稳的性能报告,应该把几件事拆开写:
| 要检查的问题 | 看哪里 | 能避免什么误判 |
|---|---|---|
| CPU 调度是否减少 | CPU trace、op 调用链 | 把 Python/dispatch 开销误认为 GPU 计算变化 |
| GPU kernel 是否变少 | CUDA timeline、kernel launch 数 | 把多个 aten op 误当多个 GPU kernel |
| GEMM 是否真的换了 | kernel 名称、tile、dtype、tn 等标记 | 把同一个 GEMM 包装成“底层计算升级” |
| 收益是否可迁移 | 不同 shape、batch、dtype 下复测 | 把单点实验写成通用结论 |
最相关的两类人,动作也不一样。
做模型推理优化的工程师,不该只给 forward 套一层 torch.compile 就交差。更应该先把 trace 分成三层:Python/CPU dispatch、CUDA kernel launch、GEMM 本体。瓶颈在哪一层,决定该调 batch、换 dtype、改图,还是继续看 Inductor 融合。
写性能报告或做迁移评估的人,也要收一下结论强度。若只在单个 Linear 上看到很小变化,就不适合推动大规模迁移。若 MLP、激活、残差、小算子链里 kernel 数明显减少,再谈 compile 收益会更稳。
接下来最该看的不是一个固定加速倍数,而是三个变量:不同 shape 下 cuBLAS 或 CUTLASS 选到的 tile 是否变化;bf16、fp16 等 dtype 是否触发不同 kernel;Inductor 对激活函数、归一化、残差结构能融合到什么程度。
这几个变量不清楚,采购新 GPU、调整 batch size、迁移到 compile,都应该慢半拍。磨刀不误砍柴工,前提是先确认刀钝在哪里。
