作为一名资深科技专栏作家，我接触过许多想要提升深度学习模型性能的用户。他们常常会采取一些“偏方”，比如使用“in-place operations”、将梯度设置为“None”、安装特定版本的PyTorch等等。
这些方法并非完全无效，但更像是炼金术而非科学。现代系统，特别是深度学习，其性能表现常常让人捉摸不透。然而，如果我们从基础原理出发，就能排除很多无效的方法，从而更高效地解决问题。
三大核心要素：计算、内存和开销
我们可以将深度学习系统的效率拆解为三个核心要素：
- 计算: GPU 用于实际浮点运算 (FLOPS) 的时间。
- 内存: 在 GPU 内部传输张量所花费的时间。
- 开销: 除此之外的一切时间消耗。
就像训练机器学习模型一样，了解系统的瓶颈所在，才能有的放矢地进行优化。例如，如果大部分时间都花在内存传输上（即内存带宽受限），那么提升 GPU 的 FLOPS 就毫无意义。反之，如果大部分时间都在进行大型矩阵乘法（即计算受限），那么用 C++ 重写模型逻辑以减少开销也无济于事。
计算：深度学习的引擎
理想情况下，我们希望最大化计算时间，毕竟我们花费了大量资金购买高性能 GPU，就应该充分利用其计算能力。然而，为了让矩阵乘法引擎高效运转，我们需要减少其他方面的耗时。
为什么 focus on 计算而不是内存带宽呢？ 因为我们无法在不改变实际操作的情况下减少所需的计算量，但可以通过优化来降低开销或内存成本。
雪上加霜的是，计算能力的增长速度远超内存带宽。下表展示了 CPU FLOPS 和内存带宽的翻倍时间：
| 指标 | 翻倍时间 |
|---|---|
| CPU FLOPS | 1 年 |
| 内存带宽 | 3 年 |
这种差距意味着，尽管 GPU 的计算能力越来越强，但如果内存带宽无法跟上，整体性能提升仍然有限。
内存带宽：数据传输的成本
内存带宽成本指的是将数据从一个地方移动到另一个地方所花费的成本。这可能包括将数据从 CPU 移动到 GPU、从一个节点移动到另一个节点，甚至从 CUDA 全局内存移动到 CUDA 共享内存。
回到工厂的比喻，GPU 的 DRAM 就好比仓库，用于存储大量数据和结果。每次执行 GPU 内核时，都需要将数据从仓库运送到工厂进行计算，然后再将结果运回仓库。
对于像 torch.cos 这样的简单操作，我们需要将数据从仓库运送到工厂，执行简单的计算，然后再将结果运回仓库。由于数据传输成本高昂，因此大部分时间都花在了数据传输上，而不是实际计算上。
操作融合：减少数据搬运的利器
为了减少内存带宽成本，我们可以采用操作融合技术。简单来说，就是将多个操作合并成一个，避免重复的数据读写。
例如，执行 x.cos().cos() 通常需要 4 次全局内存读写操作：
x1 = x.cos() # 从全局内存读取 x，写入 x1
x2 = x1.cos() # 从全局内存读取 x1，写入 x2但通过操作融合，我们只需要 2 次全局内存读写操作：
x2 = x.cos().cos() # 从全局内存读取 x，写入 x2操作融合是深度学习编译器中最重要的优化之一。它可以将多个操作合并到一起，从而节省内存带宽成本。
开销：Python 和框架的负担
开销是指代码执行过程中，除了张量传输和计算之外的所有时间消耗。例如，Python 解释器、PyTorch 框架、启动 CUDA 内核（但不执行）等都会产生开销。
现代 GPU 速度极快，而 Python 解释器却非常慢。在一个 A100 GPU 执行一次 FLOP 的时间内，Python 解释器只能执行几千万次加法运算。
PyTorch 等框架也存在多层调度机制，这也会增加开销。
为了减少开销，可以采用 JIT 编译、CUDA Graphs 等技术。
总结：对症下药，才能药到病除
总而言之，想要提升深度学习系统的性能，首先要了解系统的瓶颈所在。
| 性能瓶颈 | 解决方案 |
|---|---|
| 开销受限 | JIT 编译、操作融合、避免使用 Python |
| 内存带宽受限 | 操作融合 |
| 计算受限 | 使用 Tensor Cores、购买更强大的 GPU |
当然，用户需要考虑这些问题，本身就反映了框架设计上的不足。PyTorch 的编译器和性能分析 API 并不完善，但也在不断改进。
希望本文能够帮助你更好地理解深度学习系统的性能优化，从而让你的模型运行得更快。
参考文献
He, H. (2022). Making Deep Learning Go Brrrr From First Principles. Retrieved from https://horace.io/brrr_intro.html