大家好，我是 Guan-Ming。目前是台大電機所 CS 組的研究生，對於 AI Infra 以及 AI Compiler 的開發和研究都很有興趣。平常也有在一些開源專案擔任維護者，包含 Apache Airflow, Apache TVM 以及 PyTorch 等等。今年是我第一次參加 IT 鐵人賽，也很高興能跟「源來適你」的一群小夥伴一起參加團體賽。最近因為研究的關係碰了很多機器學習編譯器的概念以及論文，因此我想在這三十天跟大家分享以及深入的主題就是「一行 torch.compile 背後發生了什麼？30 天深度拆解 PyTorch 編譯器」！
會想寫這個系列，主要是因為 torch.compile 大概是 PyTorch 2.0 之後最常被提到、也最常被當成黑魔法的一個功能。在大部分使用 PyTorch 來建構的模型，只要加上這一行之後，模型的執行速度就會很明顯變快。不過很少人能清楚的理解它裡面到底做了些什麼？變快的原因是什麼？為什麼有時候沒效？這些問題其實都有很明確的答案，只是會需要去耐心地打開這個超大的黑盒子才看得到。所以筆者想透過這個系列跟大家一起學習、深入瞭解 PyTorch 的原始碼。並且把每一層的中間產物 dump 出來，輔以簡單的圖示或是 gif 來幫助理解。我們會一起把 PyTorch Compiler 中重要的 Dynamo、AOTAutograd、Inductor 三層，各個拆開並深入理解裡面的重要概念。最後也會有一些實作的部分，例如怎麼樣搭建一個 custom backend 等等去更理解實作中會踩到的坑。
這三十天的內容大致上會分成四個部分。第一部分是 Dynamo，講它怎麼在 Bytecode 層攔截 Python、Guard 是什麼、為什麼會 Graph Break、吐出的 FX Graph 長什麼樣。第二部分是 AOTAutograd，講 backward propagation 怎麼被一起 Trace、In-place 怎麼被正規化、Operator 怎麼被拆成基本運算。第三部分是 Inductor，講圖怎麼變成 Loop、誰跟誰融合、以及怎麼讀它生出來的 Triton 和 C++。最後一部分是整合與實戰，包括 CUDA Graph、快取、Recompilation 爆炸，以及嘗試自己寫一個 Backend。
不過第一天的目標不是馬上跳進程式碼去硬啃細節，而是先建立一個整體輪廓去理解：PyTorch 為什麼一開始選了 Eager Mode、為了拿到計算圖它試過哪些方法、為什麼那些方法沒有成功、以及最後 torch.compile 是被誰、用什麼想法做出來的。知道這段歷史，後面每一個設計決定就都有脈絡可循了。那就事不遲疑，正文準備開始！
PyTorch 的前身是 Torch7，一個用 Lua 寫的科學計算框架。2016 年，Facebook AI Research（FAIR）的 Adam Paszke、Sam Gross、Soumith Chintala、Gregory Chanan 等人把它移植到 Python，2017 年 1 月公開釋出。
當時的主流是 TensorFlow 1.x 和 Theano。它們的模式是 Define-and-Run：你先用 Python「宣告」一張靜態的計算圖，再把資料丟進 Session 執行。這種做法對編譯器很友善，因為整張圖一開始就在手上，可以隨便最佳化。但對寫程式的人非常不友善：不能在中間 print、if 要寫成 tf.cond、Debugger 也進不去。
PyTorch 反其道而行，選了 Define-by-Run，也就是我們現在說的 Eager Mode。每一行 Python 執行的當下就真的算，計算圖是在執行過程中動態被記錄下來給 Autograd 用的，用完就丟。2019 年 NeurIPS 的論文 PyTorch: An Imperative Style, High-Performance Deep Learning Library 把這個設計哲學講得很白：把易用性和「Python 就是第一公民」放在第一位，效能靠底層的 C++ 與 CUDA Kernel 撐，而不是靠犧牲彈性換來。
這個選擇讓 PyTorch 在研究圈贏了。但它也埋下了一個結構性的問題：Eager Mode 一次只看到一個 Operator，永遠不知道下一步是什麼，所以做不了任何跨 Operator 的最佳化，例如把好幾個小運算融合成一個 Kernel；也很難把模型帶離 Python，部署到別的環境。整個系列後面講的一切，其實都是在解這一個問題：在不放棄 Eager Mode 的前提下，怎麼把計算圖拿回來？
PyTorch 1.0 在 2018 年帶來了 TorchScript，提供兩條路。
第一條是 torch.jit.trace，拿一組範例輸入實際跑一次，把碰到的 Tensor 運算錄下來。這種做法很快，但它看不到控制流程，if 只會錄到當時走過的那一條分支，換一組輸入結果可能就錯了。第二條是 torch.jit.script，直接解析 Python 原始碼，翻譯成 TorchScript 自己的 IR。這樣看得到控制流程，但它只支援 Python 的一個子集，你得把模型改寫成它看得懂的樣子。
TorchScript 的目標其實偏部署，重點是把模型序列化、離開 Python 執行，加速並不是它的強項。而且它對使用者的要求太高：一碰到 dict、任意 Python 物件、第三方函式庫，就是一連串的改寫。很多團隊試過、放棄，它現在也已經進入維護模式。
2020 到 2021 年間，PyTorch 團隊在兩個方向上摸索。
一個是 torch.fx，由 James Reed 等人主導，論文 Torch.fx: Practical Program Capture and Transformation for Deep Learning in Python 發表在 MLSys 2022。它用 Python 層的 Symbolic Tracing 把 nn.Module 抓成一張很簡單的圖，並提供一套好寫的 Graph Transformation API。FX 後來成了 Dynamo 吐出的圖的格式，這點非常重要，但 FX 自己的 Tracer 一樣吃不下依賴資料的控制流程。
另一個是 Lazy Tensor，PyTorch/XLA 走的路。Tensor 運算先不真的算，累積成圖，等到有人要看結果才一次送給後端。它對使用者透明，但每一步都要重新 Trace 一次，Overhead 很高，而且一碰到需要看數值的地方就得 Flush。
這些嘗試合起來證明了同一件事：要求使用者改寫程式的方案不會贏，而能吃下任意 Python 的方案又抓不到完整的圖。要走出這個兩難，需要一個完全不同的切入點。
2021 年 9 月，Jason Ansel 在 PyTorch dev-discuss 上發表了 TorchDynamo 的雛形。它的核心想法很不一樣：不去解析 Python 原始碼，也不靠 Tracing 錄 Operator，而是利用 CPython 的 Frame Evaluation Hook（PEP 523），在 Python Bytecode 執行的當下把它攔下來，把能編的 Tensor 運算抓成 FX Graph，看不懂的地方就在那裡斷開（這就是 Graph Break），退回一般 Python 執行，之後再接回來。這樣使用者一行程式都不用改，而且永遠不會「不能跑」，最壞的情況只是沒加速而已。
Dynamo 解決的是「怎麼拿到圖」，但光有圖還不夠。同一時期還有幾塊拼圖陸續到位。AOTAutograd 由 Horace He 等人主導，它拿到前向圖之後把反向傳播也一起 Trace 出來，讓訓練也能被整張圖編譯，並且把 In-place 修改、View 這些麻煩的東西正規化成純函數式。TorchInductor 由 Jason Ansel 主導，是預設的後端，把圖 Lower 成 Loop-level IR、做融合，然後生成 Triton（GPU）或 C++（CPU）程式碼；Triton 是 OpenAI 的 Philippe Tillet 做的、用 Python 寫 GPU Kernel 的語言，讓「用 Python 生 GPU Kernel」這件事變得可行。還有 PrimTorch，把 PyTorch 兩千多個 Operator 拆解到幾百個基本 Operator，讓後端不用一一實作。
2022 年 12 月的 PyTorch Conference 上，這一整套以 PyTorch 2.0 的名義發表，2023 年 3 月正式釋出，對外的介面就是那一行 torch.compile。2024 年 ASPLOS 的論文 PyTorch 2: Faster Machine Learning Through Dynamic Python Bytecode Transformation and Graph Compilation 是這整套設計的正式論述，也是這個系列會反覆回頭看的一份文件。順帶一提，PyTorch 在 2022 年 9 月從 Meta 移交給 Linux Foundation 底下新成立的 PyTorch Foundation，但核心開發至今仍以 Meta 的團隊為主。
| 年份 | 事件 | 對「圖」的態度 |
|---|---|---|
| 2017 | PyTorch 0.1 釋出 | 純 Eager，沒有圖 |
| 2018 | PyTorch 1.0，TorchScript | 要使用者改寫成子集，才拿得到圖 |
| 2019 | NeurIPS 論文 | 明文把易用性放第一 |
| 2020 到 2021 | torch.fx、Lazy Tensor | 圖的格式與透明擷取的實驗 |
| 2021 年 9 月 | TorchDynamo 雛形 | 在 Bytecode 層攔截，抓得到就抓，抓不到就斷 |
| 2022 底到 2023 | PyTorch 2.0，torch.compile |
一行接上 Dynamo、AOTAutograd、Inductor |
| 2024 | ASPLOS 論文 | 整套設計的正式論述 |
如果把這條線縮成一句話：PyTorch 從頭到尾沒有放棄 Eager Mode，torch.compile 是在保留 Eager 語意的前提下，把圖偷偷抓出來的第三次嘗試，而前兩次的教訓決定了它的每一個設計。
回頭看這段歷史，其實可以濃縮成一個貫穿全系列的核心矛盾：易用性和完整計算圖，長期以來是二選一的。TorchScript 選了圖、犧牲了易用性，所以輸了；Lazy Tensor 選了透明，卻付出 Overhead 和 Flush 的代價。TorchDynamo 之所以是轉折點，不是因為它技術上多炫，而是它換了一個層級思考。在 Bytecode 層攔截，讓「抓不到圖」從致命錯誤降級成 Graph Break，最壞情況只是沒加速，而不是不能跑。
理解這一點很重要，因為後面會看到的每一個設計，包括 Guard 為什麼存在、Graph Break 為什麼是 feature 而不是 bug、Recompilation 為什麼會爆炸，都是「保 Eager 語意、圖能抓多少算多少」這個哲學的直接後果。帶著這個視角讀原始碼，很多看起來奇怪的取捨就會變得理所當然。
明天會從使用者的視角出發，把 torch.compile 的四段流水線攤開：Dynamo、AOTAutograd、Inductor、Runtime 各做什麼，並用 backend 參數把它們一段一段切開來親手驗證。那我們明天見！
把 Eager Mode 說成設計不是妥協，真的一下就抓到 PyTorch 的性格；先顧好「Python 就是第一公民」，再把效能交給底層撐起來，這條路走得很漂亮。從 TorchScript 卡在 trace 只能錄到當下分支、script 又得縮進子集，到 Dynamo 直接在 Bytecode 層攔截、抓不到就 Graph Break，我讀到這段很有種黑盒子終於被掀開的爽感，後面接 AOTAutograd 和 Inductor 也更有脈絡。這系列也很像把 AI 開發門檻一步步拆開來看。我手邊有多的 Lovable 額度想送給有緣人，有興趣可從連結看看我的系列。 https://ithelp.ithome.com.tw/articles/10401174
iThome鐵人賽
看影片追技術