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

/var/lib/ci-user/.local/lib/python3.10/site-packages/torch/_inductor/compile_fx.py:321: 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 = 11.43 seconds
Regional compilation time = 1.01 seconds

结论#

本指南展示了如果您的模型有重复区域，如何控制冷启动编译时间。这种方法需要用户进行修改，以便将 torch.compile 应用于重复区域，而不是更常用的全模型编译。我们正在持续致力于缩短冷启动编译时间。

脚本总运行时间： (0 分 13.489 秒)

通过局部编译缩短 torch.compile 冷启动编译时间#

先决条件#

设置#

步骤#

结论#

文档

教程

资源