标签： AI

在现代的深度学习领域，如何优化模型性能成为了一个热门话题。许多人会依赖于一些曾经有效的小技巧，比如“使用就地操作！”、“将梯度设置为 None！”、“安装 PyTorch 1.10.0 而不是 1.10.1！”等。虽然这些方法有时能带来显著的性能提升，但从基本原理出发进行分析，往往能更系统地解决问题。

理解深度学习系统的三大组成部分

深度学习系统的性能可以分解为以下三个主要组件：

1. 计算（Compute）：在 GPU 上进行实际浮点运算（FLOPS）的时间。
2. 内存（Memory）：在 GPU 内部传输张量的时间。
3. 开销（Overhead）：其他所有时间，比如 Python 解释器运行和 CUDA 内核启动等。

理解自己深度学习系统的性能瓶颈所在，可以帮助我们有针对性地进行优化。例如，如果你的系统主要花费时间在内存传输上（即内存带宽受限），那么增加 GPU 的计算能力并不会有帮助。相反，如果你主要时间都在进行大量的矩阵乘法运算（即计算受限），那么减少开销的优化也无济于事。

接下来，我们将分别讨论计算、内存带宽和开销这三个组件。

计算：如何最大化 GPU 的 FLOPS

为了充分利用 GPU 的计算能力，我们需要尽量减少在其他部分花费的时间。GPU 的 FLOPS 越高，我们的计算效率就越高。然而，计算能力增长的速度远快于内存带宽的增长速度，使得实现高效计算变得更加困难。

例如，现代机器学习加速器（如 Nvidia 的 Tensor Cores）主要针对矩阵乘法进行了优化。如果你的计算主要不是矩阵乘法，那么你将无法完全利用 GPU 的计算能力。尽管如此，其他操作（如层归一化或激活函数）所需的 FLOPS 相对较少，因此对整体性能影响不大。

内存带宽：数据传输的成本

内存带宽成本是指在系统中移动数据的成本。这包括从 CPU 到 GPU、一个节点到另一个节点，甚至从 CUDA 全局内存到 CUDA 共享内存的传输成本。频繁的数据传输会占用大量时间，导致计算资源无法充分利用。

举个例子，当我们执行一个简单的操作如 torch.cos 时，我们将数据从存储单元移到计算单元，进行计算后再将数据移回存储单元。这种频繁的数据传输非常耗时，因此我们需要尽量减少这种操作。

操作融合：减少内存开销的关键

操作融合（Operator Fusion）是深度学习编译器中最重要的优化之一。它的基本思想是将多个计算操作合并为一个，从而减少数据在全局内存和计算单元之间的传输次数。例如，执行 x.cos().cos() 时，通常需要进行四次全局读取和写入，而通过操作融合，我们只需两次全局内存读取和写入。

这种优化在执行大规模计算时尤其重要，可以显著减少内存带宽成本，提高整体性能。

开销：代码执行之外的时间

开销指的是代码在执行实际计算和传输数据之外所花费的时间。现代 GPU 的计算速度非常快，而 Python 解释器相对较慢，因此在 Python 解释器中花费的时间会导致 GPU 资源闲置。

如何识别和减少开销

要识别是否受到开销限制，可以通过增加数据量来观察运行时间的变化。如果增加数据量后运行时间没有成比例增加，那么系统可能受到开销限制。使用 PyTorch 的分析工具也可以帮助识别 CPU 和 GPU 之间的配合情况，从而找到优化方向。

结论

要提升深度学习系统的性能，关键是识别模型的性能瓶颈，并有针对性地进行优化。以下是不同性能瓶颈的优化策略：

性能瓶颈	可能的解决方案
开销受限	使用 tracing，操作融合，避免使用 Python，采用真正的 JIT
内存带宽受限	操作融合
计算受限	使用 Tensor Cores，升级 GPU 硬件

理解这些基本原理和优化策略，可以帮助我们更有效地提升深度学习模型的性能。

---

参考文献：

• Horace He, Making Deep Learning Go Brrrr From First Principles

作为一名资深科技专栏作家，我接触过许多想要提升深度学习模型性能的用户。他们常常会采取一些“偏方”，比如使用“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✅