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 階段，雖然對大多數主流模型支援良好，但對於包含極複雜動態分支或冷門自定義運算元的專案，依然可能出現編譯失敗或退化回傳統模式的情況。我們建議在生產部署前，先在小批次樣本上跑一輪數值一致性校驗。