torch-compile

torch.compile：用 TorchDynamo 捕获计算图，再由 Inductor 生成融合 kernel

补充说明

torch.compile 在运行时捕获 Python 代码中的张量计算图，交给 TorchInductor 编译成更高效的 kernel。它能融合逐元素操作、减少 Python 开销、生成 Triton/C++ kernel，常用于训练和推理加速。它不是“万能加速器”：动态控制流、图断裂、不支持的算子、频繁变化的 shape 都会削弱收益。

常见模式：default 平衡编译时间和性能，reduce-overhead 降低 Python/启动开销（适合小模型/推理），max-autotune 花更多时间搜索 kernel（适合稳定负载）。先用小步数 warmup，再测稳定后的吞吐，不要把首次编译时间算进性能。

常用示例

import torch

model = MyModel().cuda()
model = torch.compile(model)

# 训练循环：先 warmup 几步，再开始计时
for i, batch in enumerate(loader):
    if i == 5:
        torch.cuda.synchronize()
        start = torch.cuda.Event(enable_timing=True)
        start.record()
    loss = model(batch).sum()
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

# 不同模式
model = torch.compile(model, mode="reduce-overhead")
model = torch.compile(model, mode="max-autotune")

# 动态 shape：减少因 shape 变化导致的重复编译
model = torch.compile(model, dynamic=True)

# 强制整图编译，遇到图断裂直接报错，便于定位
model = torch.compile(model, fullgraph=True)

# 查看图断裂原因
TORCH_LOGS=graph_breaks python train.py

# 查看编译/重编译统计
TORCH_LOGS=recompiles python train.py

易错点与排错

• 把首次编译时间算进性能：第一次运行包含捕获和编译，往往很慢；先 warmup，再测稳定状态。
• 动态 shape 导致反复重编译：batch size、序列长度变化会触发 recompile；尽量固定 shape，或用 dynamic=True 并接受一定性能损失。
• 图断裂：Python 控制流、.item()、numpy 调用、打印、动态索引都会打断图；用 TORCH_LOGS=graph_breaks 定位。
• fullgraph=True 直接报错：说明模型有断点；先去掉该选项定位，再逐步修复或接受 graph break。
• 与 FSDP/DeepSpeed 的顺序不兼容：不同版本对 compile 与分布式包装的顺序要求不同；用小模型验证，并看官方示例。
• 训练精度变化：融合/自动混合精度可能改变数值行为；用固定种子和评测集对比 loss/指标。
• 编译缓存占磁盘：Inductor 缓存会增长；CI/容器中设置缓存目录并定期清理。
• 不是所有算子都更快：小 tensor、memory-bound 或已高度优化的算子收益有限；用 profiler 对比，不要盲目全模型 compile。

相关命令

• fsdp — 分布式分片训练。
• nsys — 看编译后的时间线与 kernel。
• ncu — 分析生成的 kernel。
• torchrun — 多卡启动。
• liger-kernel — 另一类融合 kernel 优化。

参考链接

• torch.compile 文档
• TorchInductor
• torch.compile 教程

读完自测

合上页面，先回答，再展开对照。

1. torch.compile 的 default、reduce-overhead、max-autotune 分别适合什么场景？
2. 什么是 graph break？怎样定位？
3. 为什么动态 shape 会导致性能下降？