PyTorch torch.compile与Dynamo原理
请解释PyTorch 2.0的torch.compile。Dynamo/AOTAutograd/Inductor三者的分工是什么？使用torch.compile可以带来多少加速？什么场景下不适合使用？
回答
我是大山
torch.compile是PyTorch 2.0的JIT编译器。
三层架构：
- Dynamo：捕获Python图，将PyTorch代码转为FX Graph
- AOTAutograd：生成反向传播图
- Inductor：将计算图编译为GPU Kernel(Triton)或CPU代码
使用：
@torch.compile
model = torch.compile(model, mode="reduce-overhead")
加速效果：
- 训练：5-50%(取决于模型)
- 推理：10-200%
- Transformer类模型收益最大
不适用场景：
- 动态控制流复杂
- 小模型/batch size小(编译开销大于收益)
- 有大量Python side effect的代码