PyTorch 2.0 正式发布：引入 torch.compile，模型加速无需改代码

PyTorch 2.0 是主流深度学习框架 PyTorch 的大版本更新。它在保留大家熟悉的 Python 动态图开发习惯的同时，彻底重构了底层编译器体系，让你不用修改现有代码逻辑，就能直接给模型训练和推理加速。

这次更新凝聚了自 1.13.1 版本以来的 4541 次提交和 428 位贡献者。核心变化是引入了全新的编译体系，同时完全向前兼容。

它能做什么与怎么做到的

新版本最核心的功能是引入了 torch.compile 这个可选 API。你只需要用它把原有模型包裹一层，系统就会返回一个编译优化后的新模型。

这一功能底层由四个核心组件驱动。TorchDynamo 负责利用 Python 栈帧评估钩子安全捕获动态计算图，AOTAutograd 负责提前生成反向传播追踪，PrimTorch 把两千多个复杂的 PyTorch 算子精简成约 250 个基础算子，最后的 TorchInductor 则是深度学习编译器。

在英伟达和 AMD GPU 上，TorchInductor 深度集成了 OpenAI Triton 编译器，能自动生成媲美手写 CUDA 内核的高性能代码。官方在 165 个开源模型上进行了测试，torch.compile 成功支持了 93% 的模型，在 float32 精度下平均提速 20%，在 AMP 混合精度下平均提速 36%。

除了通用编译，针对当前最火的 Transformer 架构，PyTorch 2.0 把缩放点积注意力算子直接内置到了 torch.nn.functional 模块中。它原生集成了来自 HazyResearch 的 FlashAttention 以及 xFormers 的高效注意力实现，你再也不需要手动折腾复杂的第三方依赖。

硬件生态方面，苹果 Mac 平台的 MPS 后端扩充了高频算子支持，覆盖算子数量超过 300 个，跑 Stable Diffusion 和 Whisper 会更稳定。AWS 则针对 Graviton3 架构的 C7g 实例优化了 CPU 推理，常见模型如 ResNet-50 和 BERT 的推理速度都有明显提升。

怎么跑起来

体验 PyTorch 2.0 非常简单，通常只需要更新 pip 或 conda 包即可安装对应环境。官方提供了详尽的 2.0 快速入门文档。

在代码层面，最基础的调用方式就是在现有模型定义后加一行代码，写成 model = torch.compile(model)。原本的训练循环和前向推理逻辑完全不需要调整。

如果你想在 Mac 上跑 GPU 加速，建议升级到最新的 macOS 版本，MPS 后端的兼容性与正确性测试在最新系统下表现最好。

为什么重要

对于大部分正在训练大语言模型或做日常模型微调的开发者来说，PyTorch 2.0 绝对值得立刻尝试。过去想要获得极速的训练体验，通常要面对 JAX 这类静态图框架更高的学习成本，或者四处寻找第三方的 CUDA 算子补丁。现在只要套上原生一行代码就能吃到编译红利，研发门槛大幅降低。

不过编译功能目前仍处于 Beta 阶段，虽然对大多数主流模型支持良好，但对于包含极复杂动态分支或冷门自定义算子的项目，依然可能出现编译失败或退化回传统模式的情况。我们建议在生产部署前，先在小批次样本上跑一轮数值一致性校验。