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

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

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

該工作同時于系統(tǒng)領(lǐng)域重要國際會議 USENIX ATC’24 發(fā)表長文，第一作者清華大學(xué)博士生張晨、通訊作者為翟季冬教授。PACMAN 實驗室在機(jī)器學(xué)習(xí)系統(tǒng)領(lǐng)域持續(xù)深入研究。MagPy 是繼 PET、EINNET 等工作后在深度學(xué)習(xí)編譯器上的又一次探索。欲了解更多相關(guān)成果可查看翟季冬教授首頁：https://pacman.cs.tsinghua.edu.cn/~zjd

• 論文地址：https://www.usenix.org/system/files/atc24-zhang-chen.pdf
• 項目地址：https://github.com/heheda12345/MagPy

研究背景：深度學(xué)習(xí)計算圖提取技術(shù)

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

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

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

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

MagPy 的解決方案

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

MagPy 的設(shè)計基于以下幾點觀察：

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

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

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

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

實驗

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

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