推理加速 --- torch.compile

一、背景介绍

PyTorch 2.0 官宣了一个重要特性 —— torch.compile，这一特性将 PyTorch 的性能推向了新的高度，并将 PyTorch 的部分内容从 C++ 移回 Python。torch.compile 是一个完全附加的（可选的）特性，因此 PyTorch 2.0 是 100% 向后兼容的。

支撑 torch.compile 的技术包括研发团队新推出的 TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor。

• TorchDynamo：使用 Python Frame Evaluation Hooks 安全地捕获 PyTorch 程序，这项重大创新是 PyTorch 过去 5 年来在安全图结构捕获方面的研发成果汇总；
• AOTAutograd：重载 PyTorch 的 autograd 引擎，作为一个跟踪 autodiff，用于生成 ahead-of-time 向后跟踪；
• PrimTorch：将约 2000 多个 PyTorch 算子归纳为一组约 250 个原始算子的闭集，开发人员可以将其作为构建完整 PyTorch 后端的目标。这大大降低了编写 PyTorch 功能或后端的流程；
• TorchInductor：是一种深度学习编译器，可为多个加速器和后端生成快速代码。对于 NVIDIA GPU，它使用 OpenAI Triton 作为关键构建块。
• TorchDynamo、AOTAutograd、PrimTorch 和 TorchInductor 是用 Python 编写的，并支持 dynamic shapes（无需重新编译就能发送不同大小的向量），这使得它们具备灵活、易于破解的特性，降低了开发人员和供应商的使用门槛。

1.2 首次推理速度慢

值得注意的是，torch.compile函数不会立即对函数进行加速优化，而是在第一次运行的时候才进行优化。这就会导致编译后的函数在第一次执行的时候十分缓慢。但是，会在第二次执行开始，变得非常快。

1.3 推理多次之后，又会出现一次速度特别慢的情况，感觉好像是重新优化

• 原因分析

动态推理的原因，即输入的tensor是动态的

二、如何使用

一行代码进行使用

modoel_compile = torch.compile(model)

三、其他设置

def compile(model: Optional[Callable] = None, *,
             fullgraph: builtins.bool = False,
             dynamic: builtins.bool = False,
             backend: Union[str, Callable] = "inductor",
             mode: Union[str, None] = None,
             options: Optional[Dict[str, Union[str, builtins.int, builtins.bool]]] = None,
             disable: builtins.bool = False) -> Callable:
     """
    Optimizes given model/function using TorchDynamo and specified backend.
    Args:
        model (Callable): Module/function to optimize
        fullgraph (bool): Whether it is ok to break model into several subgraphs
        dynamic (bool): Use dynamic shape tracing
        backend (str or Callable): backend to be used
        mode (str): Can be either "default", "reduce-overhead" or "max-autotune"
        options (dict): A dictionary of options to pass to the backend.
        disable (bool): Turn torch.compile() into a no-op for testing
    """

3.1 mode 设置

关于mode参数介绍如下：

• （默认）default: 适合加速大模型，编译速度快且无需额外存储空间
• reduce-overhead：适合加速小模型，需要额外存储空间
• max-autotune：编译速度非常耗时，但提供最快的加速

• 未正确设置该参数导致的问题：
  1）推理速度慢

• 示例

modoel_compile = torch.compile(model， mode="reduce-overhead")

3.2 backend

backend 编译器后端:API使用哪个后端将中间表示(IR)计算图(FX graph)转换为低级内核操作。这个选项对于调试graph编译问题和更好地理解torch.compile的内部非常有用。在大多数情况下，默认的Inductor后端似乎能够提供最佳的训练性能结果。有很多后端列表，我们可以使用下面命令查看：

from torch import _dynamo
 print(_dynamo.list_backends())

我们测试使用nvprims-nvfuser后端，可以获得比eager模式13%的性能提升(与默认后端28.6%的性能提升相比)。具体区别还是要看Pytorch文档，我们这里就不细说了，因为文档都有。

3.3 fullgraph

fullgraph 强制单个图:这个参数是非常有用，可以确保没有任何不希望的图截断。

3.4 dynamic

dynamic 动态形状:目前 2.0对具有动态形状的张量的编译支持在某种程度上是有限的。编译具有动态形状的模型的一个常见解决方案是重新编译，但会大大增加开销并大大降低训练速度。如果您的模型确实包含动态形状，将动态标志设置为True将带来更好的性能，特别是减少重新编译的次数。

都有什么是动态形状呢，最简单的就是时间序列或文本长度不同，如果不进行对齐操作的话序列长度不同就是动态的形状。

参考资料

https://www.zhihu.com/question/590338541/answer/2959785643
https://blog.csdn.net/tMb8Z9Vdm66wH68VX1/article/details/130939176