开始¶
让我们从一个简单的示例开始。请注意，您可能会看到更多 GPU 越新，速度就越快。
import torch
def fn(x, y):
a = torch.cos(x).cuda()
b = torch.sin(y).cuda()
return a + b
new_fn = torch.compile(fn, backend="inductor")
input_tensor = torch.randn(10000).to(device="cuda:0")
a = new_fn(input_tensor, input_tensor)
此示例实际上不会运行得更快。其目的是演示
和 功能是
逐点运算的示例，例如它们在
向量。您可能想要使用的更著名的逐点运算是
是类似 .在 Eager 模式下的逐点作是
次优，因为每个 API 都需要从
memory 进行一些更改，然后写回这些更改。单
Inductor 所做的最重要的优化是 fusion。那么回到我们的
例如，我们可以将 2 次读取和 2 次写入转换为 1 次读取和 1 次写入，其中
尤其对于瓶颈是内存的较新 GPU 至关重要
bandwidth （将数据发送到 GPU 的速度）而不是 compute
（您的 GPU 处理浮点运算的速度）。torch.cos()torch.sin()torch.relu()
感应器提供的另一个主要优化是自动的 支持 CUDA 图形。 CUDA 图形有助于消除启动单个 内核，这与较新的 GPU 特别相关。
TorchDynamo 支持许多不同的后端，但 inductor 特别有效
通过生成 Triton 内核和
我们可以通过使用实际生成的内核来检查它们TORCH_COMPILE_DEBUG=1 python trig.py
@pointwise(size_hints=[16384], filename=__file__, meta={'signature': {0: '*fp32', 1: '*fp32', 2: 'i32'}, 'device': 0, 'constants': {}, 'configs': [instance_descriptor(divisible_by_16=(0, 1, 2), equal_to_1=())]})
@triton.jit
def kernel(in_ptr0, out_ptr0, xnumel, XBLOCK : tl.constexpr):
xnumel = 10000
xoffset = tl.program_id(0) * XBLOCK
xindex = xoffset + tl.reshape(tl.arange(0, XBLOCK), [XBLOCK])
xmask = xindex < xnumel
x0 = xindex
tmp0 = tl.load(in_ptr0 + (x0), xmask)
tmp1 = tl.sin(tmp0)
tmp2 = tl.sin(tmp1)
tl.store(out_ptr0 + (x0 + tl.zeros([XBLOCK], tl.int32)), tmp2, xmask)
您可以验证两者的融合是否确实发生了
因为这两个作发生在单个 Triton 内核中
临时变量保存在访问速度非常快的 registers 中。sinsin
您可以在此处阅读有关 Triton 性能的更多信息，但关键是它是用 Python 编写的 因此，即使您没有写完所有这些，您也可以轻松理解它 许多 CUDA 内核。
接下来，让我们尝试一个真实的模型，比如 PyTorch 中的 resnet50 枢纽。
import torch
model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
opt_model = torch.compile(model, backend="inductor")
model(torch.randn(1,3,64,64))
这不是唯一可用的后端，您可以在 REPL 中运行以查看所有可用的后端。尝试 or next 作为灵感。torch._dynamo.list_backends()cudagraphsnvfuser
现在让我们做一些更有趣的事情，我们的社区经常 使用来自 transformer 或 TIMM 的预训练模型以及 我们的设计目标是让 Dynamo 和电感器开箱即用 人们想要创作的任何模型。
因此，我们将直接从 HuggingFace 中心并对其进行优化：
import torch
from transformers import BertTokenizer, BertModel
# Copy pasted from here https://huggingface.co/bert-base-uncased
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained("bert-base-uncased").to(device="cuda:0")
model = torch.compile(model, backend="inductor") # This is the only line of code that we changed
text = "Replace me by any text you'd like."
encoded_input = tokenizer(text, return_tensors='pt').to(device="cuda:0")
output = model(**encoded_input)
如果从模型中删除 和 ，则 Triton 将生成 C++ 内核，该内核将
针对在 CPU 上运行进行了优化。您可以检查 Triton 或 C++
kernels 进行 BERT 的 Bert 计算，它们显然比三角函数更复杂
我们上面有示例，但您可以同样略读它并了解您是否
了解 PyTorch。to(device="cuda:0")encoded_input
同样，让我们尝试一个 TIMM 示例
import timm
import torch._dynamo as dynamo
import torch
model = timm.create_model('resnext101_32x8d', pretrained=True, num_classes=2)
opt_model = torch.compile(model, backend="inductor")
opt_model(torch.randn(64,3,7,7))
我们对 Dynamo 和 inductor 的目标是构建覆盖率最高的 ML 编译器 这应该适用于您扔给它的任何模型。
现有后端¶
TorchDynamo 的后端列表不断增长，可以在 backends 文件夹中找到
或者每个 API 都有其可选的依赖项。torch._dynamo.list_backends()
一些最常用的后端包括：
- 训练和推理后端：
- 仅推理后端：
为什么需要另一种优化 PyTorch 代码的方法？¶
虽然 PyTorch 中存在许多其他代码优化工具 ecosystem，它们都有自己的 flow。 以下是现有方法及其限制的几个示例：
torch.jit.trace()如果无法跟踪，则为 silently 错误，例如： 在控制流期间torch.jit.script()需要修改用户代码或库代码 通过添加类型注释并删除非 PyTorch 代码torch.fx.symbolic_trace()要么正确跟踪，要么给出一个硬 错误，但它仅限于可跟踪代码，因此仍然无法处理 控制流torch._dynamo开箱即用并生成部分图形。 它仍然可以选择生成某些人所需的单个图表 情况，但允许更平滑的过渡，其中部分图形可以 无需修改代码即可优化nopython=True