深度学习模型跑得慢,很多团队的第一反应仍是翻经验帖:改 in-place、调 batch size、换 PyTorch 版本、试一个新编译器。PyTorch 团队成员 Horace He 在 9 月发布的文章《Making Deep Learning go Brrrr From First Principles》提醒,性能优化不该从技巧清单开始,而应先判断系统到底卡在计算、显存带宽,还是框架和调度开销。

这篇文章的重要性不在于给出一条万能加速公式。它真正有用的地方,是把训练和推理性能拆回三个区间:compute-bound、memory-bound 和 overhead-bound。不同区间对应不同解法,错配优化方向,往往只是在昂贵 GPU 上做无效劳动。

GPU 优化的目标,是把时间尽量留给矩阵乘

现代 AI 加速器的峰值算力,主要来自矩阵乘专用硬件。以 NVIDIA GPU 为例,Tensor Core 承担了深度学习中最重的矩阵乘任务;原文提到的 A100 规格也说明,标称峰值 FLOPS 并不等于任意算子都能达到的性能。非矩阵算子能用到的通用计算能力要小得多。

这解释了一个工程上常见的错觉:LayerNorm、激活函数、逐元素操作的 FLOPS 占比很小,却可能花掉不成比例的时间。原因通常不是它们“算得慢”,而是它们频繁把中间结果写回全局显存,再读回来。计算单元在等数据,GPU 就不会真正“跑满”。

瓶颈类型典型现象更匹配的优化方向错配动作
计算受限大矩阵乘占主导,achieved FLOPS 接近峰值提升 matmul 形状、精度、Tensor Core 利用率只减少 kernel launch
显存带宽受限带宽接近上限,FLOPS 利用率低算子融合、重计算、减少中间读写只换更高 FLOPS GPU
调度开销受限小算子多,GPU 空转明显编译、批量化、减少 kernel launch手写复杂数学优化

这也是为什么算子融合会成为深度学习编译器的核心能力。NVFuser、XLA、TorchInductor、Triton 这类工具解决的不是“把数学变少”,而是尽量避免中间张量在全局显存里来回搬运。对一串逐元素操作来说,融合后读一次、写一次,往往比每个算子各自读写更接近硬件真实能力。

算子融合有效,但不是线性加速按钮

原文用 x.cos().cos() 解释了融合的直觉:不融合时,中间结果要写回显存再读出;融合后,可以在同一个 kernel 里连续完成计算。这个例子能说明为什么很多激活函数的实际耗时差距没有数学表达式看起来那么大,因为瓶颈常常在搬数据,不在多做几个标量运算。

但工程上不能把“融合”理解成所有场景的保险药。它的收益取决于是否真的减少了全局显存访问,也取决于算子形态、张量大小、寄存器压力、并行度和编译器能否生成足够好的 kernel。融合过度也可能让单个 kernel 变复杂,反而降低占用率。

这里有一个行业现实:公开宣传里,GPU 常被拿峰值 FLOPS 做横向比较;实际优化时,工程师更该盯 achieved FLOPS / peak FLOPS、显存带宽利用率和 compute intensity。前者告诉你算力吃到了多少,后者告诉你每搬一个字节做了多少计算。只看标称算力,很容易把 memory-bound 问题误判成“硬件不够强”。

对工程团队的影响:少试玄学,多做定位

受影响最大的是训练平台团队、推理服务团队和需要自己写 CUDA/Triton kernel 的算法工程师。预算有限时,买更贵的 GPU 未必比改 kernel、上编译器或调整算子结构更划算;延迟敏感的推理服务里,小算子过多带来的 launch 开销,也可能比单个算子的数学量更致命。

更实用的判断顺序是:先用 profiler 看时间分布,再用 FLOPS、带宽和 kernel launch 数量定位区间。若带宽接近上限而 FLOPS 很低,优先看融合、重计算、减少中间张量读写;若矩阵乘占主要耗时,重点检查维度是否适配 Tensor Core、精度路径是否正确;若 GPU 利用率断断续续,编译、批量化和减少小 kernel 才是正路。

接下来最该观察的,不是哪一个“神奇参数”又被社区验证有效,而是 PyTorch 2.x 编译栈、Triton kernel 和厂商库能否把更多真实模型自动带到正确区间。自动编译器很强,但模型结构、动态 shape、分支逻辑和自定义算子仍会让优化结果变得不稳定。性能工程没有捷径,只有先量后改。