减少 torch.compile 冷启动编译时间：区域编译方法

# For tips on running notebooks in Google Colab, see
# https://pytorch.org/tutorials/beginner/colab
%matplotlib inline

减少 torch.compile 冷启动编译时间：区域编译方法#

作者: Animesh Jain

随着深度学习模型的规模越来越大，这些模型的编译时间也随之增加。这种延长的编译时间会导致推理服务中的启动时间变长，或者在大规模训练中浪费资源。本教程展示了一个示例，说明如何通过选择编译模型中的重复区域而不是整个模型来减少冷启动编译时间。

先决条件#

  • Pytorch 2.5 或更高版本

设置#

在我们开始之前，如果尚未安装 torch，我们需要先进行安装。

pip install torch
注意:

此功能从 2.5 版本开始提供。如果您使用的是 2.4 版本，可以通过启用配置标志 torch._dynamo.config.inline_inbuilt_nn_modules=True 来防止在区域编译期间重新编译。在 2.5 版本中，此标志默认启用。

from time import perf_counter

步骤#

在本教程中，我们将按照以下步骤进行：

  1. 导入所有必要的库。

  2. 定义并初始化具有重复区域的神经网络。

  3. 理解全模型编译与区域编译之间的区别。

  4. 测量全模型编译和区域编译的编译时间。

首先，让我们导入加载数据所需的必要库：

import torch
import torch.nn as nn

接下来，让定义并初始化具有重复区域的神经网络。

通常，神经网络由重复的层组成。例如，大型语言模型由许多 Transformer 块组成。在本教程中，将使用 nn.Module 类创建 Layer，作为重复区域的代理。然后，将创建由 64 个 Layer 类实例组成的 Model。

class Layer(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear1 = torch.nn.Linear(10, 10)
        self.relu1 = torch.nn.ReLU()
        self.linear2 = torch.nn.Linear(10, 10)
        self.relu2 = torch.nn.ReLU()

    def forward(self, x):
        a = self.linear1(x)
        a = self.relu1(a)
        a = torch.sigmoid(a)
        b = self.linear2(a)
        b = self.relu2(b)
        return b


class Model(torch.nn.Module):
    def __init__(self, apply_regional_compilation):
        super().__init__()
        self.linear = torch.nn.Linear(10, 10)
        # Apply compile only to the repeated layers.
        if apply_regional_compilation:
            self.layers = torch.nn.ModuleList(
                [torch.compile(Layer()) for _ in range(64)]
            )
        else:
            self.layers = torch.nn.ModuleList([Layer() for _ in range(64)])

    def forward(self, x):
        # In regional compilation, the self.linear is outside of the scope of `torch.compile`.
        x = self.linear(x)
        for layer in self.layers:
            x = layer(x)
        return x

接下来，回顾一下全模型编译与区域编译之间的区别。

在全模型编译中，整个模型作为整体进行编译。这是大多数用户在使用 torch.compile 时的常见做法。在这个示例中，将 torch.compile 应用于 Model 对象。这将有效地内联 64 个层，生成需要编译的大型图。您可以通过运行此教程并设置 TORCH_LOGS=graph_code 来查看完整的图。

model = Model(apply_regional_compilation=False).cuda()
full_compiled_model = torch.compile(model)

另一方面，区域编译则是编译模型的一部分。通过有策略地选择编译模型的重复区域，可以编译更小的图，然后为所有区域重用这个编译后的图。在示例中，torch.compile 仅应用于 layers，而不是整个模型。

regional_compiled_model = Model(apply_regional_compilation=True).cuda()

将编译应用于重复区域而不是整个模型，可以大幅节省编译时间。在这里，只需编译 Layer 实例，然后在 Model 对象中重复使用它 64 次。

请注意，对于重复区域，模型的一部分可能不会被编译。例如，Model 中的 self.linear 就不在区域编译的范围内。

此外，需要注意的是，性能加速与编译时间之间存在权衡。全模型编译涉及更大的图，理论上提供了更多的优化空间。然而，在实际应用中，根据模型的不同，我们观察到许多情况下全模型编译与区域编译之间的加速差异很小。

接下来，测量全模型编译和区域编译的编译时间。

torch.compile 是即时（JIT）编译器，这意味着它在第一次调用时进行编译。在下面的代码中，测量第一次调用所花费的总时间。虽然这种方法并不精确，但它提供了很好的估计，因为大部分时间都花在编译上。

def measure_latency(fn, input):
    # Reset the compiler caches to ensure no reuse between different runs
    torch.compiler.reset()
    with torch._inductor.utils.fresh_inductor_cache():
        start = perf_counter()
        fn(input)
        torch.cuda.synchronize()
        end = perf_counter()
        return end - start


input = torch.randn(10, 10, device="cuda")
full_model_compilation_latency = measure_latency(full_compiled_model, input)
print(f"Full model compilation time = {full_model_compilation_latency:.2f} seconds")

regional_compilation_latency = measure_latency(regional_compiled_model, input)
print(f"Regional compilation time = {regional_compilation_latency:.2f} seconds")

assert regional_compilation_latency < full_model_compilation_latency

/media/pc/data/lxw/envs/anaconda3a/envs/ai/lib/python3.12/site-packages/onnxscript/converter.py:820: FutureWarning: 'onnxscript.values.Op.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
  param_schemas = callee.param_schemas()
/media/pc/data/lxw/envs/anaconda3a/envs/ai/lib/python3.12/site-packages/onnxscript/converter.py:820: FutureWarning: 'onnxscript.values.OnnxFunction.param_schemas' is deprecated in version 0.1 and will be removed in the future. Please use '.op_signature' instead.
  param_schemas = callee.param_schemas()
/media/pc/data/lxw/envs/anaconda3a/envs/ai/lib/python3.12/site-packages/torch/_inductor/compile_fx.py:167: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
  warnings.warn(

Full model compilation time = 35.90 seconds
Regional compilation time = 2.04 seconds

结论#

本教程展示了如何在模型具有重复区域时控制冷启动编译时间。这种方法需要用户修改代码，将 torch.compile 应用于重复区域，而不是通常使用的全模型编译。