Python程式到計算圖一鍵轉化，詳解清華開源深度學習編譯器MagPy
机器之心發表於2024-10-06
AIxiv專欄是機器之心釋出學術、技術內容的欄目。過去數年，機器之心AIxiv專欄接收報導了2000多篇內容，覆蓋全球各大高校與企業的頂級實驗室，有效促進了學術交流與傳播。如果您有優秀的工作想要分享，歡迎投稿或者聯絡報導。投稿郵箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com
張晨，清華大學計算機系高效能所博士生，導師為翟季冬老師，主要研究方向為面向人工智慧和量子計算的高效能異構計算系統。在OSDI、SC、ATC、ICS會議上發表一作論文，並獲得 ICS21 最佳學生論文。曾獲得 SC19, SC20, ISC21 國際超級計算機競賽冠軍。獲清華大學本科生特等獎學金、國家獎學金、北京市優秀畢業生、北京市優秀畢業設計等榮譽。2024 年 7 月，清華大學計算機系 PACMAN 實驗室釋出開源深度學習編譯器 MagPy，可一鍵編譯使用者使用 Python 編寫的深度學習程式，實現模型的自動加速。儘管目前存在大量高效能的深度學習編譯器，但是這些編譯器均以計算圖作為輸入，需要由使用者將編寫的 Python 程式手動轉化為計算圖。為了避免這種不便性，該團隊設計了 MagPy，直接面向使用者編寫的 Python+PyTorch 程式，自動將其轉化為適用於深度學習編譯器的計算圖表示，從而充分發揮深度學習編譯器的最佳化能力，避免使用者使用複雜 Python 語法帶來的效能下降，為使用者帶來易用性和效率的雙豐收。該工作同時於系統領域重要國際會議 USENIX ATC’24 發表長文，第一作者清華大學博士生張晨、通訊作者為翟季冬教授。PACMAN 實驗室在機器學習系統領域持續深入研究。MagPy 是繼 PET、EINNET 等工作後在深度學習編譯器上的又一次探索。欲瞭解更多相關成果可檢視翟季冬教授首頁：https://pacman.cs.tsinghua.edu.cn/~zjd- 論文地址：https://www.usenix.org/system/files/atc24-zhang-chen.pdf
- 專案地址：https://github.com/heheda12345/MagPy
近年來，深度學習在生物科學、天氣預報和推薦系統等多個領域展示了其強大能力。為了簡化程式設計過程，使用者傾向於使用 Python 編寫深度學習模型，並在需要進行張量操作時呼叫如 PyTorch 等的張量庫。此時，使用者程式會在呼叫張量庫時立即執行張量操作，如此不加最佳化地直接執行程式效能較差。另一方面，為了提升深度學習模型的執行速度，深度學習編譯器傾向於使用以運算元圖的格式表示的深度學習模型作為輸入，在計算圖上進行圖級最佳化，如圖替換和運算元融合。當可以獲取到模型的計算圖時，代表性的深度學習編譯器 TorchInductor 和 XLA 可以在 PyTorch 的基礎上平均加速模型 1.47 倍和 1.40 倍。具體結果如圖 1 所示，標記為 Fullgraph-Inductor 和 Fullgraph-XLA。然而，實現這種加速的前提是使用者手動將程式轉換為計算圖格式，這對許多模型開發者來說是困難的。尤其是隨著深度學習的廣泛應用，越來越多的模型是由化學、生物和天文學等領域的非專業程式設計師開發的。因此，迫切需要一種自動化方法將使用者編寫的 Python 程式轉換為編譯器友好的圖格式來加速程式，這被稱為計算圖提取技術。由於 Python 程式具有極強的動態性，加之使用者程式存在行為的不確定性，現有的計算圖提取技術在處理較複雜的使用者程式時無法取得最優的效能，如圖 1 中的 TorchDynamo-Inductor（使用 TorchDynamo 提取計算圖，使用 TorchInductor 編譯）、 LazyTensor-XLA（使用 LazyTensor 追蹤計算圖，使用 XLA 編譯）所示。圖 1 ：深度學習編譯器可以顯著提升模型執行效率，但現有的圖提取技術阻礙了這一點。圖中 Eager 表示直接執行 PyTorch 程式，Fullgraph-Inductor 與 Fullgraph-XLA 分別表示 Inductor、XLA 對模型的計算圖進行編譯後的加速，TorchDynamo-Inductor 與 LazyTensor-XLA 分別表示使用 TorchDynamo 和 LazyTensor 技術從使用者 Python 程式中提取計算圖再進行編譯的效能。MagPy 的核心思想是分析 Python 直譯器中的執行狀態資訊，從而讓編譯器能夠更好的理解使用者程式。Python 直譯器能夠準確支援所有 Python 特性，並在執行時保留了高層次的執行狀態資訊，如各個變數的型別和值等等。透過有效利用直譯器提供的資訊，能夠更全面地瞭解程式的行為，從而更好地提取程式計算圖。首先，大多數深度學習程式的動態性是有限的。儘管這些程式是用 Python 編寫的，具有資料型別、控制流邏輯和執行時函式排程等潛在的動態特性，但其計算圖結構在不同批次間通常保持不變。ParityBench 是一個從 Github 上自動爬取超過 100 顆星的 PyTorch 深度學習程式組成的基準測試集，它的 1421 個程式中，83% 的程式（1191 個）均滿足有限動態性的性質。對於這些程式，可以透過在程式執行過程中監控張量操作，較為簡便地獲取其計算圖。根據這個性質，MagPy 將計算圖提取問題從分析 “計算圖是什麼” 簡化為分析 “得到的計算圖何時會發生變化”。其次，只有外部值能影響程式行為。利用這一特性，可以更簡易地檢測出會導致計算圖發生變化的因素。這裡的 “程式行為” 包括計算圖的結構和所有程式副作用（side effect）。只要程式從外部讀取的所有值（如輸入引數和全域性變數）保持不變，且呼叫的函式的輸出結果不具有隨機性，程式行為就不會發生變化。因此，MagPy 只需驗證所有從外部讀取的值都不變，即可保證計算圖結構不變。例如，儘管圖 2 中的程式使用了許多複雜的 Python 特性，但只要所有從外部讀取的值（如 x、dims、self.scale 和 self.dim，標記為粗體）與之前執行一致，計算圖就保持不變。MagPy 會首先執行一個 “守衛函式” 對於這些值是否發生變化進行檢查（Guards），當檢查透過時，MagPy 將會執行一個 “模擬函式”(mock code)，用以呼叫經過深度學習編譯器編譯的計算圖及模擬程式的所有副作用（如示例中的對 x 進行賦值）。第三，守衛函式和模擬函式都可以透過分析程式執行狀態來確定。守衛函式的作用是驗證新一次執行的輸入狀態是否與之前執行匹配，模擬函式的目的是重現之前執行的最終狀態。這兩個部分僅基於執行時狀態，而不是使用者程式的邏輯。Python 直譯器在解釋執行程式的過程中，保留了所有需要的執行狀態資訊，因此不再需要具體分析 Python 複雜而動態的執行邏輯。守衛函式和模擬函式需要關注的變數包括顯式讀取或寫入外部的值（如 self）以及被它們引用的值（如 self.dim）。因此，MagPy 設計了引用關係圖來記錄和分析程式行為。基於上述觀察，MagPy 提出了引用關係圖（Reference Graph，簡寫為 RefGraph）來記錄程式執行期間的程式狀態。MagPy 定義了執行狀態介面，用於在程式執行期間收集執行時資訊，並使用基於標註的圖更新規則來維護 RefGraph。MagPy 還提出了在 RefGraph 上進行遍歷生成守衛函式和模擬函式的演算法。具體細節可以閱讀論文。MagPy 具有極高的 Python 語言特性覆蓋率，其在對 ParityBench 中 1191 個靜態的真實使用者程式進行測試時，成功將 93.40% 的程式轉化為完整的運算子圖，大幅高於現有工作 TorchScript（35%）和 TorchDynamo（77.2%）由於更完整的計算圖匯出，MagPy 在端到端測試中，也表現出具有競爭力的效能。下圖展示了對於影像處理、自然語言處理等典型深度學習模型，MagPy 取得的加速。MagPy 可取得最高 2.88 倍，平均 1.55 倍的提升。實驗在單張 A100 上進行，X-Y 表示使用圖匯出技術 X 和圖層編譯器 Y。