跳转到主要内容
博客

赋能模型性能：通用模型转换方法的艺术

引言

PyTorch 2.0 (PT2) 提供了一个编译执行模式，它重写 Python 字节码以提取 PyTorch 操作序列，并将其转换为图IR。然后，IR 通过可定制的后端进行即时编译，无需用户干预即可提高训练性能。通常，生产模型可能需要经过多个优化/下沉阶段才能达到性能目标。因此，拥有一个编译模式是理想的，因为它可以将提高模型性能的工作与直接修改 PyTorch 模型实现的工作分开。因此，编译模式变得更加重要，使 PyTorch 用户无需修改 PyTorch 代码实现即可提高模型性能。此功能对于优化复杂模型，包括大规模和生产就绪模型，尤其有价值。

在我们之前的博文中，我们概述了如何使用启发式模型转换规则来优化复杂的生产模型。虽然这些规则使一些试点模型获得了显著的性能提升，但它们缺乏普遍的适应性；它们不能在不同的模型中，甚至有时在单个模型的不同部分中始终表现良好。

Fig.1 PT1 Graph mode vs PT2 Compile mode.

图1：PT1 图模式 vs PT2 编译模式。

在这篇博文中，我们提出了一种更通用的模型转换解决方案，作为 PT2 编译器的插件，如图1所示，它更通用、性能更好、更用户友好，无需人工干预即可提高模型训练和推理的性能。如图2所示，通过将以前用户定义的转换纳入编译器，我们简化了生产堆栈。这些更改为更广泛的 PyTorch 模型带来了优势，而不仅仅是 Meta 模型，这已经集成到 PT2 中，并可供所有 PyTorch 模型使用。

Fig.2 Simplified stack with PT2 compile mode.

图2：PT2 编译模式简化了堆栈。

指导原则：原子规则

传统上，人们可能会使用预定义的启发式规则来用另一个性能更好的子图替换模型子图，以减少启动开销，最小化内存带宽，并充分占用 SM。然而，这种方法扩展性不佳，因为很难制定一套完美适用于所有模型的规则。

与其纠缠于庞大复杂的规则，我们实际上可以把它们分解成更小、更容易理解的部分——我们称之为“原子规则”。这些微小的效率利器旨在转换单个算子，进行融合/转换的一个步骤。这使得它们易于处理和应用，为优化模型提供了一条简单的路径。因此，有了这些原子规则，优化任何模型以获得顶级性能变得轻而易举！

我们将通过一些简单的例子来演示我们如何使用一系列原子规则来替换复杂的启发式规则。

案例1：以访问嵌入表开始的计算链的水平融合

水平融合意味着将并行算子融合为一个，以减少要启动的内核数量并提高性能。在我们之前的博客（第3.2节）中，我们描述了在嵌入包之后融合 LayerNorm 和激活函数的模型转换，如图所示。然而，这种方法有局限性：

  1. 它只适用于嵌入后的 LayerNorm 和激活函数。
  2. 它仅限于具有特定架构规则的模型，导致我们的生产堆栈出现各种问题，包括参数更改和推理中断。

为了改进，我们可以使用图3所示的三个原子规则来替换复杂的启发式规则：

  • 水平融合同一拆分节点后的 LayerNorm。
  • 然后，水平融合同一拆分节点后的 tanh 函数。
  • 最后，融合垂直的拆分-拼接节点。

这些原子规则为模型简化和优化提供了一种清晰、精简的方法。

Fig.3 Before, we optimized the model in one go by replacing subgraphs. Now, with atomic rules, we optimize step-by-step, covering more cases.

图3：以前，我们通过替换子图一步到位地优化模型。现在，使用原子规则，我们逐步优化，涵盖更多情况。

案例2：融合水平 MLP

MLP（多层感知器）是深度神经网络的基本组成部分，通常由线性、归一化和激活函数组成。在复杂模型中，通常需要融合许多水平 MLP。传统方法会找到并用融合模块替换并行 MLP，如图4所示，但这并非总是直截了当。有些模型可能没有归一化，或者它们可能使用不同的激活函数，这使得应用一刀切的规则变得困难。

这就是我们的原子规则派上用场的地方。这些简化的规则一次针对一个单独的算子，使过程更容易、更易于管理。我们使用以下原子规则进行水平 MLP 融合：

  • 融合水平线性算子
  • 融合水平 LayerNorm。
  • 融合水平激活函数。
Fig.4 Pseudocode for fusing MLP. Traditional optimizations need manual Python code changes.

图4：融合 MLP 的伪代码。传统优化需要手动修改 Python 代码。

这些规则的优点在于它们不仅限于一种情况。它们可以广泛应用。由于 PyTorch 模型是用 torch 算子构建的，因此关注一小组算子可以简化过程。与编写特定的庞大模式替换规则相比，这种方法不仅更易于管理，而且更通用，使其更容易有效地优化各种模型。

我们的原则是使用链式原子规则来替换启发式规则。虽然这种方法涵盖了更广泛的情况，但它确实需要更长的图搜索和模式匹配时间。下一个问题是：我们如何在高效执行编译时图搜索的同时最小化编译时间？

我们设计了一个两步贪婪算法，如图5所示。此过程的第一步是识别目标节点，我们遵循某些规则，例如，识别所有具有相同输入形状的线性操作。一旦识别，我们使用广度优先搜索 (BFS) 策略将这些节点分成不同的集合，以便集合中的节点没有数据依赖性。每个集合中的节点都是独立的，可以水平融合。

Fig.5 Process of model transformation with graph IR.

图5：使用图 IR 进行模型转换的过程。

使用我们的方法，对于我们最大的内部模型之一，搜索时间大约为60秒，这对于即时任务来说是可管理的。

总结

在我们对内部排名模型的测试中，我们观察到在 torch.compile 带来的性能提升的基础上，五个模型的训练性能平均提高了 5% 到 15%。我们已经在 PT2 编译器堆栈中启用了优化，并在用户选择 Inductor 作为后端时将其默认启用（配置）。我们期望我们通用的转换方法可以使 Meta 之外的模型受益，并期待通过这种编译器级别转换框架进行更多讨论和改进。

致谢

非常感谢 Mark Saroufim、Gregory Chanan、Adnan Aziz 和 Rocky Liu 提供的详细而富有洞察力的评论。