通过区域编译减少 torch.compile 冷启动编译时间

创建日期：2024 年 10 月 10 日 | 最后更新日期：2024 年 10 月 16 日 | 最后验证日期：2024 年 10 月 10 日

作者： 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 作为重复区域的代理。然后，我们将创建一个由该 Layer 类的 64 个实例组成的 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()

将编译应用于重复区域而不是完整模型，可以显著节省编译时间。在这里，我们将只编译一个层实例，然后在 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

Full model compilation time = 24.83 seconds
Regional compilation time = 1.10 seconds

结论

本秘籍展示了如果您的模型包含重复区域，如何控制冷启动编译时间。这种方法需要用户修改，将 torch.compile 应用于重复区域，而不是更常用的完整模型编译。我们正在持续努力减少冷启动编译时间。