torch.compile 端到端教程"

intermediate/torch_compile_full_example

在 Google Colab 中运行

Colab

下载 Notebook

Notebook

在 GitHub 上查看

GitHub

torch.compile 端到端教程

作者: William Wen

torch.compile 是加速 PyTorch 代码的新方式！torch.compile 通过将 PyTorch 代码实时（JIT）编译为优化后的内核，使 PyTorch 代码运行得更快，且几乎不需要修改代码。

本教程涵盖了一个使用 torch.compile 训练和评估真实模型的端到端示例。如需了解 torch.compile 的入门介绍，请查看 torch.compile 介绍教程。

必需的 pip 依赖

• torch >= 2.0

• torchvision

您将学到什么

• 如何将 torch.compile 应用于真实模型

• torch.compile 在真实模型上的加速效果

• 由于编译开销，torch.compile 的前几次迭代通常较慢

先决条件

• torch.compile 简介

# NOTE: a modern NVIDIA GPU (H100, A100, or V100) is recommended for this tutorial in
# order to reproduce the speedup numbers shown below and documented elsewhere.

import torch
import warnings

gpu_ok = False
if torch.cuda.is_available():
    device_cap = torch.cuda.get_device_capability()
    if device_cap in ((7, 0), (8, 0), (9, 0)):
        gpu_ok = True

if not gpu_ok:
    warnings.warn(
        "GPU is not NVIDIA V100, A100, or H100. Speedup numbers may be lower "
        "than expected."
    )

/var/lib/workspace/intermediate_source/torch_compile_full_example.py:51: UserWarning: GPU is not NVIDIA V100, A100, or H100. Speedup numbers may be lower than expected.
  warnings.warn(

让我们演示使用 torch.compile 如何加速真实模型。我们将通过在随机数据上评估和训练一个 torchvision 模型，来比较标准 Eager 模式和 torch.compile 的性能。

开始之前，我们需要定义一些实用函数。

# Returns the result of running `fn()` and the time it took for `fn()` to run,
# in seconds. We use CUDA events and synchronization for the most accurate
# measurements.
def timed(fn):
    start = torch.cuda.Event(enable_timing=True)
    end = torch.cuda.Event(enable_timing=True)
    start.record()
    result = fn()
    end.record()
    torch.cuda.synchronize()
    return result, start.elapsed_time(end) / 1000


# Generates random input and targets data for the model, where `b` is
# batch size.
def generate_data(b):
    return (
        torch.randn(b, 3, 128, 128).cuda(),
        torch.randint(1000, (b,)).cuda(),
    )


N_ITERS = 10

from torchvision.models import densenet121


def init_model():
    return densenet121().cuda()

首先，让我们比较推理性能。

注意，在调用 torch.compile 时，我们使用了额外的 mode 参数，我们将在下文讨论它。

model = init_model()

# Note that we generally recommend directly compiling a torch.nn.Module by calling
# its .compile() method.
model_opt = init_model()
model_opt.compile(mode="reduce-overhead")

inp = generate_data(16)[0]
with torch.no_grad():
    print("eager:", timed(lambda: model(inp))[1])
    print("compile:", timed(lambda: model_opt(inp))[1])

eager: 0.3411640319824219
/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(
compile: 55.4854921875

请注意，与 Eager 模式相比，torch.compile 完成任务所需的时间更长。这是因为 torch.compile 在执行时将模型编译为优化后的内核。在我们的示例中，模型结构没有改变，因此不需要重新编译。所以，如果我们多运行几次优化后的模型，应该能看到相比 Eager 模式的显著改进。

eager_times = []
for i in range(N_ITERS):
    inp = generate_data(16)[0]
    with torch.no_grad():
        _, eager_time = timed(lambda: model(inp))
    eager_times.append(eager_time)
    print(f"eager eval time {i}: {eager_time}")

print("~" * 10)

compile_times = []
for i in range(N_ITERS):
    inp = generate_data(16)[0]
    with torch.no_grad():
        _, compile_time = timed(lambda: model_opt(inp))
    compile_times.append(compile_time)
    print(f"compile eval time {i}: {compile_time}")
print("~" * 10)

import numpy as np

eager_med = np.median(eager_times)
compile_med = np.median(compile_times)
speedup = eager_med / compile_med
assert speedup > 1
print(
    f"(eval) eager median: {eager_med}, compile median: {compile_med}, speedup: {speedup}x"
)
print("~" * 10)

eager eval time 0: 0.01965363121032715
eager eval time 1: 0.018131967544555663
eager eval time 2: 0.01760767936706543
eager eval time 3: 0.01725644874572754
eager eval time 4: 0.01682329559326172
eager eval time 5: 0.01767731285095215
eager eval time 6: 0.017535999298095704
eager eval time 7: 0.01746329689025879
eager eval time 8: 0.017618976593017577
eager eval time 9: 0.016737279891967775
~~~~~~~~~~
compile eval time 0: 0.0783062744140625
compile eval time 1: 0.008770560264587402
compile eval time 2: 0.009492480278015136
compile eval time 3: 0.008492032051086425
compile eval time 4: 0.008455167770385743
compile eval time 5: 0.008292351722717285
compile eval time 6: 0.00818995189666748
compile eval time 7: 0.008111104011535645
compile eval time 8: 0.008136704444885253
compile eval time 9: 0.008073216438293456
~~~~~~~~~~
(eval) eager median: 0.017571839332580566, compile median: 0.008373759746551513, speedup: 2.098440827588469x
~~~~~~~~~~

确实，我们可以看到使用 torch.compile 运行模型带来了显著的加速。加速主要源于减少 Python 开销和 GPU 读/写操作，因此观察到的加速效果可能会因模型架构和批大小等因素而异。例如，如果模型架构很简单且数据量很大，那么瓶颈将是 GPU 计算，此时观察到的加速效果可能不那么明显。

根据选择的 mode 参数，你可能会看到不同的加速结果。"reduce-overhead" 模式使用 CUDA 图来进一步减少 Python 的开销。对于你自己的模型，你可能需要尝试不同的模式以实现最大加速。你可以点击此处阅读更多关于模式的内容。

你可能还会注意到，我们第二次使用 torch.compile 运行模型时比其他运行要慢得多，尽管它仍然比第一次运行快得多。这是因为 "reduce-overhead" 模式会为 CUDA 图运行一些预热迭代。

现在，让我们考虑比较训练性能。

model = init_model()
opt = torch.optim.Adam(model.parameters())


def train(mod, data):
    opt.zero_grad(True)
    pred = mod(data[0])
    loss = torch.nn.CrossEntropyLoss()(pred, data[1])
    loss.backward()
    opt.step()


eager_times = []
for i in range(N_ITERS):
    inp = generate_data(16)
    _, eager_time = timed(lambda: train(model, inp))
    eager_times.append(eager_time)
    print(f"eager train time {i}: {eager_time}")
print("~" * 10)

model = init_model()
opt = torch.optim.Adam(model.parameters())

# Note that because we are compiling a regular Python function, we do not
# call any .compile() method.
train_opt = torch.compile(train, mode="reduce-overhead")

compile_times = []
for i in range(N_ITERS):
    inp = generate_data(16)
    _, compile_time = timed(lambda: train_opt(model, inp))
    compile_times.append(compile_time)
    print(f"compile train time {i}: {compile_time}")
print("~" * 10)

eager_med = np.median(eager_times)
compile_med = np.median(compile_times)
speedup = eager_med / compile_med
assert speedup > 1
print(
    f"(train) eager median: {eager_med}, compile median: {compile_med}, speedup: {speedup}x"
)
print("~" * 10)

eager train time 0: 0.34415924072265625
eager train time 1: 0.05326233673095703
eager train time 2: 0.05141708755493164
eager train time 3: 0.051241985321044924
eager train time 4: 0.0516577262878418
eager train time 5: 0.050902015686035154
eager train time 6: 0.050098175048828124
eager train time 7: 0.0510750732421875
eager train time 8: 0.050132991790771485
eager train time 9: 0.05106175994873047
~~~~~~~~~~
compile train time 0: 166.327453125
compile train time 1: 2.359182373046875
compile train time 2: 0.022718463897705078
compile train time 3: 0.02105855941772461
compile train time 4: 0.02045952033996582
compile train time 5: 0.020478975296020507
compile train time 6: 0.020700159072875975
compile train time 7: 0.020564992904663085
compile train time 8: 0.020479999542236327
compile train time 9: 0.02040115165710449
~~~~~~~~~~
(train) eager median: 0.05115852928161621, compile median: 0.020632575988769532, speedup: 2.479502768314639x
~~~~~~~~~~

同样，我们可以看到 torch.compile 在第一次迭代中花费的时间更长，因为它必须编译模型，但在随后的迭代中，与 Eager 模式相比，我们看到了显著的加速。

我们在此强调，本教程中提供的加速数值仅供演示参考。官方加速值可以在 TorchInductor 性能仪表板上查看。

结论

在本教程中，我们将 torch.compile 应用于真实模型的训练和推理，并展示了加速效果。

重要的是，我们要指出，由于编译开销，编译后的模型在前几次迭代中比 Eager 模式慢，但预计在随后的迭代中会有加速。

如需了解 torch.compile 的入门介绍，请查看 torch.compile 介绍教程。

如需排查问题并深入了解如何将 torch.compile 应用于你的代码，请查看 torch.compile 编程模型。

希望你能尝试使用 torch.compile！