張晨，清華大學計算機系高性能所博士生，導師為翟季冬老師，主要研究方向為面向人工智能和量子計算的高性能異構計算系統(tǒng)。在OSDI、SC、ATC、ICS會議上發(fā)表一作論文，并獲得 ICS21 最佳學生論文。曾獲得 SC19, SC20, ISC21 國際超級計算機競賽冠軍。獲清華大學本科生特等獎學金、國家獎學金、北京市優(yōu)秀畢業(yè)生、北京市優(yōu)秀畢業(yè)設計等榮譽。

2024 年 7 月，清華大學計算機系 PACMAN 實驗室發(fā)布開源深度學習編譯器 MagPy，可一鍵編譯用戶使用 Python 編寫的深度學習程序，實現(xiàn)模型的自動加速。

盡管目前存在大量高性能的深度學習編譯器，但是這些編譯器均以計算圖作為輸入，需要由用戶將編寫的 Python 程序手動轉化為計算圖。為了避免這種不便性，該團隊設計了 MagPy，直接面向用戶編寫的 Python+PyTorch 程序，自動將其轉化為適用于深度學習編譯器的計算圖表示，從而充分發(fā)揮深度學習編譯器的優(yōu)化能力，避免用戶使用復雜 Python 語法帶來的性能下降，為用戶帶來易用性和效率的雙豐收。

該工作同時于系統(tǒng)領域重要國際會議 USENIX ATC’24 發(fā)表長文，第一作者清華大學博士生張晨、通訊作者為翟季冬教授。PACMAN 實驗室在機器學習系統(tǒng)領域持續(xù)深入研究。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

研究背景：深度學習計算圖提取技術

近年來，深度學習在生物科學、天氣預報和推薦系統(tǒng)等多個領域展示了其強大能力。為了簡化編程過程，用戶傾向于使用 Python 編寫深度學習模型，并在需要進行張量操作時調用如 PyTorch 等的張量庫。此時，用戶程序會在調用張量庫時立即執(zhí)行張量操作，如此不加優(yōu)化地直接執(zhí)行程序性能較差。另一方面，為了提升深度學習模型的運行速度，深度學習編譯器傾向于使用以算子圖的格式表示的深度學習模型作為輸入，在計算圖上進行圖級優(yōu)化，如圖替換和算子融合。當可以獲取到模型的計算圖時，代表性的深度學習編譯器 TorchInductor 和 XLA 可以在 PyTorch 的基礎上平均加速模型 1.47 倍和 1.40 倍。

具體結果如圖 1 所示，標記為 Fullgraph-Inductor 和 Fullgraph-XLA。然而，實現(xiàn)這種加速的前提是用戶手動將程序轉換為計算圖格式，這對許多模型開發(fā)者來說是困難的。尤其是隨著深度學習的廣泛應用，越來越多的模型是由化學、生物和天文學等領域的非專業(yè)程序員開發(fā)的。因此，迫切需要一種自動化方法將用戶編寫的 Python 程序轉換為編譯器友好的圖格式來加速程序，這被稱為計算圖提取技術。

由于 Python 程序具有極強的動態(tài)性，加之用戶程序存在行為的不確定性，現(xiàn)有的計算圖提取技術在處理較復雜的用戶程序時無法取得最優(yōu)的性能，如圖 1 中的 TorchDynamo-Inductor（使用 TorchDynamo 提取計算圖，使用 TorchInductor 編譯）、 LazyTensor-XLA（使用 LazyTensor 追蹤計算圖，使用 XLA 編譯）所示。

圖 1 ：深度學習編譯器可以顯著提升模型運行效率，但現(xiàn)有的圖提取技術阻礙了這一點。圖中 Eager 表示直接執(zhí)行 PyTorch 程序，F(xiàn)ullgraph-Inductor 與 Fullgraph-XLA 分別表示 Inductor、XLA 對模型的計算圖進行編譯后的加速，TorchDynamo-Inductor 與 LazyTensor-XLA 分別表示使用 TorchDynamo 和 LazyTensor 技術從用戶 Python 程序中提取計算圖再進行編譯的性能。

MagPy 的解決方案

MagPy 的核心思想是分析 Python 解釋器中的執(zhí)行狀態(tài)信息，從而讓編譯器能夠更好的理解用戶程序。Python 解釋器能夠準確支持所有 Python 特性，并在運行時保留了高層次的執(zhí)行狀態(tài)信息，如各個變量的類型和值等等。通過有效利用解釋器提供的信息，能夠更全面地了解程序的行為，從而更好地提取程序計算圖。

MagPy 的設計基于以下幾點觀察：

首先，大多數深度學習程序的動態(tài)性是有限的。盡管這些程序是用 Python 編寫的，具有數據類型、控制流邏輯和運行時函數調度等潛在的動態(tài)特性，但其計算圖結構在不同批次間通常保持不變。ParityBench 是一個從 Github 上自動爬取超過 100 顆星的 PyTorch 深度學習程序組成的基準測試集，它的 1421 個程序中，83% 的程序（1191 個）均滿足有限動態(tài)性的性質。對于這些程序，可以通過在程序執(zhí)行過程中監(jiān)控張量操作，較為簡便地獲取其計算圖。根據這個性質，MagPy 將計算圖提取問題從分析 “計算圖是什么” 簡化為分析 “得到的計算圖何時會發(fā)生變化”。

其次，只有外部值能影響程序行為。利用這一特性，可以更簡易地檢測出會導致計算圖發(fā)生變化的因素。這里的 “程序行為” 包括計算圖的結構和所有程序副作用（side effect）。只要程序從外部讀取的所有值（如輸入參數和全局變量）保持不變，且調用的函數的輸出結果不具有隨機性，程序行為就不會發(fā)生變化。因此，MagPy 只需驗證所有從外部讀取的值都不變，即可保證計算圖結構不變。例如，盡管圖 2 中的程序使用了許多復雜的 Python 特性，但只要所有從外部讀取的值（如 x、dims、self.scale 和 self.dim，標記為粗體）與之前運行一致，計算圖就保持不變。MagPy 會首先運行一個 “守衛(wèi)函數” 對于這些值是否發(fā)生變化進行檢查（Guards），當檢查通過時，MagPy 將會運行一個 “模擬函數”(mock code)，用以調用經過深度學習編譯器編譯的計算圖及模擬程序的所有副作用（如示例中的對 x 進行賦值）。

第三，守衛(wèi)函數和模擬函數都可以通過分析程序執(zhí)行狀態(tài)來確定。守衛(wèi)函數的作用是驗證新一次執(zhí)行的輸入狀態(tài)是否與之前運行匹配，模擬函數的目的是重現(xiàn)之前運行的最終狀態(tài)。這兩個部分僅基于運行時狀態(tài)，而不是用戶程序的邏輯。Python 解釋器在解釋執(zhí)行程序的過程中，保留了所有需要的執(zhí)行狀態(tài)信息，因此不再需要具體分析 Python 復雜而動態(tài)的執(zhí)行邏輯。守衛(wèi)函數和模擬函數需要關注的變量包括顯式讀取或寫入外部的值（如 self）以及被它們引用的值（如 self.dim）。因此，MagPy 設計了引用關系圖來記錄和分析程序行為。

基于上述觀察，MagPy 提出了引用關系圖（Reference Graph，簡寫為 RefGraph）來記錄程序執(zhí)行期間的程序狀態(tài)。MagPy 定義了執(zhí)行狀態(tài)接口，用于在程序執(zhí)行期間收集運行時信息，并使用基于標注的圖更新規(guī)則來維護 RefGraph。MagPy 還提出了在 RefGraph 上進行遍歷生成守衛(wèi)函數和模擬函數的算法。具體細節(jié)可以閱讀論文。

實驗

MagPy 具有極高的 Python 語言特性覆蓋率，其在對 ParityBench 中 1191 個靜態(tài)的真實用戶程序進行測試時，成功將 93.40% 的程序轉化為完整的操作符圖，大幅高于現(xiàn)有工作 TorchScript（35%）和 TorchDynamo（77.2%）

由于更完整的計算圖導出，MagPy 在端到端測試中，也表現(xiàn)出具有競爭力的性能。下圖展示了對于圖像處理、自然語言處理等典型深度學習模型，MagPy 取得的加速。MagPy 可取得最高 2.88 倍，平均 1.55 倍的提升。實驗在單張 A100 上進行，X-Y 表示使用圖導出技術 X 和圖層編譯器 Y。