開發者狂喜！Meta最新發布的LLM Compiler，實現77%自動調優效率
机器之心發表於2024-07-01
Meta搞了個很牛的LLM Compiler，幫助程式設計師更高效地寫程式碼。
昨天，三大 AI 巨頭 OpenAI、Google、Meta 組團釋出自家大模型最新研究成果 ——OpenAI 推出基於 GPT-4 訓練的專門找 bug 的新模型 CriticGPT，谷歌開源 9B、27B 版 Gemma2，而 Meta 則拿出了一項最新的人工智慧突破 ——LLM Compiler。這是一套強大的開源模型，旨在最佳化程式碼並徹底改變編譯器設計。這項創新有可能改變開發者處理程式碼最佳化的方式，使其更快、更高效、更經濟。據悉，該 LLM Compiler 的最佳化潛力達到了自動調優搜尋的 77%，這一結果可以顯著減少編譯時間，並提高各種應用的程式碼效率，並且在反彙編方面，其往返反彙編的成功率為 45%。有網友表示，這像是程式碼最佳化和反彙編的遊戲規則改變者。大語言模型在眾多軟體工程和程式設計任務中表現出了卓越的能力，然而它們在程式碼最佳化及編譯器領域的應用尚未得到充分挖掘。訓練這些 LLMs 需要消耗大量的計算資源，包括昂貴的 GPU 時間和龐大的資料集，這往往使得許多研究和專案難以為繼。為了彌補這一空白，Meta 研究團隊引入了一種 LLM Compiler，以專門最佳化程式碼並徹底改變編譯器設計。透過在包含 5460 億個標記的 LLVM-IR 和彙編程式碼的龐大語料庫上訓練模型，他們使模型能夠理解編譯器中間表示、組合語言和最佳化技術。論文連結：https://ai.meta.com/research/publications/meta-large-language-model-compiler-foundation-models-of-compiler-optimization/研究人員在他們的論文中解釋說：「LLM Compiler 增強了對編譯器中間表示（IR）、組合語言和最佳化技術的理解。」這種增強的理解使模型能夠執行以前僅限於人類專家或專業工具的任務。LLM Compiler 的訓練流程如圖 1 所示。LLM Compiler 在程式碼大小最佳化方面取得了顯著成果。在測試中，模型的最佳化潛力達到了自動調優搜尋的 77%，這一結果可以顯著減少編譯時間，並提高各種應用的程式碼效率。模型在反彙編方面的能力更為出色。LLM Compiler 在將 x86_64 和 ARM 彙編程式碼轉換回 LLVM-IR 時，往返反彙編的成功率為 45%（其中 14% 完全匹配）。這種能力對於逆向工程任務和舊程式碼維護可能具有無法估量的價值。專案的核心貢獻者之一 Chris Cummins 強調了這項技術的潛在影響：「透過提供兩種大小的預訓練模型（7 億和 13 億引數）並透過微調版本展示其有效性，」他說，「LLM Compiler 為探索 LLM 在程式碼和編譯器最佳化領域未被開發的潛力鋪平了道路。」用於訓練程式設計 LLMs 的資料通常主要由像 Python 這樣的高階源語言組成，彙編程式碼在這些資料集中的比例微乎其微，編譯器 IR 的比例更小。為了構建一個對這些語言有良好理解的 LLM，研究團隊用 Code Llama 的權重初始化 LLM Compiler 模型，然後在一個以編譯器為中心的資料集上訓練 4010 億個 token，這個資料集主要由彙編程式碼和編譯器 IR 組成，如表 1 所示。資料集 LLM Compiler 主要在由 LLVM（版本 17.0.6）生成的編譯器中間表示和彙編程式碼上進行訓練，這些資料來源於用於訓練 Code Llama 的同一資料集，已在表 2 中概述了該資料集。與 Code Llama 一樣，我們也從自然語言資料集中獲取少量訓練批次。
為了理解程式碼最佳化的機制，研究團隊對 LLM Compiler 模型進行指令微調，以模擬編譯器最佳化，如圖 2 所示。其思路是從有限的未最佳化種子程式集合中，透過對這些程式應用隨機生成的編譯器最佳化序列，生成大量示例。然後他們訓練模型預測最佳化生成的程式碼，還訓練模型預測應用最佳化後的程式碼大小。任務規範。給定未經最佳化的 LLVM-IR（由 clang 前端輸出），一個最佳化過程列表，以及一個起始程式碼大小，生成應用這些最佳化後的結果程式碼以及程式碼大小。這個任務有兩種型別：在第一種中，模型預期輸出編譯器 IR；在第二種中，模型預期輸出彙編程式碼。兩種型別的輸入 IR、最佳化過程和程式碼大小是相同的，提示決定了所需的輸出格式。程式碼大小。他們使用兩個指標來衡量程式碼大小：IR 指令數和二進位制大小。二進位制大小透過將 IR 或彙編降級為目標檔案後，.TEXT 和 .DATA 段大小的總和計算得出。我們排除 .BSS 段，因為它不影響磁碟上的大小。最佳化 pass。在這項工作中，研究團隊針對 LLVM 17.0.6，並使用新的過程管理器 (PM, 2021)，它將 pass 分類為不同的級別，如模組、函式、迴圈等，以及轉換和分析 pass 。轉換 pass 改變給定的輸入 IR，而分析 pass 生成影響後續轉換的資訊。在 opt 的 346 個可能的 pass 引數中，他們選擇了 167 個使用。這包括每個預設最佳化流水線 (例如 module (default<Oz>))，單獨的最佳化轉換 pass (例如 module (constmerge))，但排除了非最佳化實用程式 pass (例如 module (dot-callgraph)) 和不保留語義的轉換 pass (例如 module (internalize))。他們排除了分析 pass，因為它們沒有副作用，我們依賴 pass 管理器根據需要注入依賴的分析 pass。對於接受引數的 pass，我們使用預設值 (例如 module (licm<allowspeculation>))。表 9 包含了所有使用的 pass 列表。我們使用 LLVM 的 opt 工具應用 pass 列表，並使用 clang 將結果 IR 降級為目標檔案。清單 1 顯示了使用的命令。資料集。研究團隊透過對錶 2 中總結的未最佳化程式應用 1 到 50 個隨機最佳化 pass 列表生成編譯器模擬資料集。每個 pass 列表的長度是均勻隨機選擇的。pass 列表是透過從上述 167 個 pass 集合中均勻取樣生成的。導致編譯器崩潰或在 120 秒後超時的 pass 列表被排除。LLM Compiler FTD ：擴充套件下游編譯任務 操作編譯器標誌對執行時效能和程式碼大小都有顯著影響。研究團隊訓練 LLM Compiler FTD 模型執行下游任務，即為 LLVM 的 IR 最佳化工具 opt 選擇標誌，以生成最小的程式碼大小。標誌調優的機器學習方法以前已經顯示出良好的結果，但在不同程式之間的泛化方面存在困難。以前的工作通常需要編譯新程式數十或數百次，以嘗試不同的配置並找出效能最佳的選項。該研究團隊透過預測標誌來最小化未見程式的程式碼大小，在這個任務的零樣本版本上訓練和評估 LLM Compiler FTD 模型。他們的方法不依賴於所選擇的編譯器和最佳化指標，他們打算在未來針對執行時效能。目前，最佳化程式碼大小簡化了訓練資料的收集。任務規範。研究團隊向 LLM Compiler FTD 模型呈現一個未最佳化的 LLVM-IR (由 clang 前端生成)，並要求它生成應該應用的 opt 標誌列表，這些最佳化應用前後的二進位制大小，以及輸出程式碼，如果無法對輸入程式碼進行改進，則生成一個只包含未最佳化二進位制大小的簡短輸出訊息。他們使用了與編譯器模擬任務相同的受限最佳化 pass 集，並以相同的方式計算二進位制大小。圖 3 說明了用於生成訓練資料的過程以及如何在推理時使用模型。在評估時只需要生成的 pass 列表。他們從模型輸出中提取 pass 列表，並使用給定的引數執行 opt。然後，研究人員可以評估模型預測的二進位制大小和最佳化輸出程式碼的準確性，但這些是輔助學習任務，不是使用所必需的。正確性。LLVM 最佳化器並非無懈可擊，以意外或未經測試的順序執行最佳化 pass 可能會暴露出微妙的正確性錯誤，從而降低模型的實用性。為了緩解這種風險，研究團隊開發了 PassListEval，這是一個工具，用於幫助自動識別破壞程式語義或導致編譯器崩潰的 pass 列表。圖 4 顯示了該工具的概覽。PassListEval 接受候選 pass 列表作為輸入，並在一個包含 164 個自測試 C++ 程式的套件上對其進行評估，這些程式取自 HumanEval-X。每個程式都包含一個程式設計挑戰的參考解決方案，例如「檢查給定數字向量中是否有兩個數字之間的距離小於給定閾值」，以及驗證正確性的單元測試套件。他們將候選 pass 列表應用於參考解決方案，然後將它們與測試套件連結以生成二進位制檔案。執行時，如果任何測試失敗，二進位制檔案將崩潰。如果任何二進位制崩潰，或者任何編譯器呼叫失敗，我們就拒絕該候選 pass 列表。資料集。該團隊在一個源自 450 萬個未最佳化 IR 的標誌調優示例資料集上訓練了 LLM Compiler FTD 模型，這些 IR 用於預訓練。為生成每個程式的最佳 pass 列表示例，他們進行了廣泛的迭代編譯過程，如圖 3 所示。1. 研究團隊使用大規模隨機搜尋為程式生成初始候選最佳 pass 列表。對每個程式，他們獨立生成最多 50 個 pass 的隨機列表，從之前描述的 167 個可搜尋 pass 集合中均勻取樣。每次他們評估一個程式的 pass 列表時，都記錄生成的二進位制大小，然後選擇產生最小二進位制大小的每個程式 pass 列表。他們執行了 220 億次獨立編譯，平均每個程式 4,877 次。2. 隨機搜尋生成的 pass 列表可能包含冗餘 pass，這些 pass 對最終結果沒有影響。此外，一些 pass 順序是可交換的，重新排序不會影響最終結果。由於這些會在訓練資料中引入噪聲，他們開發了一個最小化過程，並將其應用於每個 pass 列表。最小化包括三個步驟：冗餘 pass 消除、氣泡排序和插入搜尋。在冗餘 pass 消除中，他們透過迭代刪除單個 pass 來最小化最佳 pass 列表，看它們是否對二進位制大小有貢獻，如果沒有，就丟棄它們。重複此過程，直到不能再丟棄 pass。然後氣泡排序嘗試為 pass 子序列提供統一排序，根據關鍵字對 pass 進行排序。最後，插入排序透過遍歷 pass 列表中的每個 pass 並嘗試在其之前插入 167 個搜尋 pass 中的每一個來執行區域性搜尋。如果這樣做改善了二進位制大小，就保留這個新的 pass 列表。整個最小化管道迴圈直到達到固定點。最小化後的 pass 列表長度分佈如圖 9 所示。平均 pass 列表長度為 3.84。3. 他們將之前描述過 PassListEval 應用於候選最佳 pass 列表。透過這種方式，他們確定了 1,704,443 個獨立 pass 列表中的 167,971 個 (9.85%) 會導致編譯時或執行時錯4. 他們將 100 個最常見的最優 pass 列表廣播到所有程式，如果發現改進就更新每個程式的最佳 pass 列表。之後，唯一最佳 pass 列表的總數從 1,536,472 減少到 581,076。上述自動調優管道相比 -Oz 產生了 7.1% 的幾何平均二進位制大小減少。圖 10 顯示了單個 pass 的頻率。對他們來說，這種自動調優作為每個程式最佳化的黃金標準。雖然發現的二進位制大小節省很顯著，但這需要 280 億次額外編譯，計算成本超過 21,000 個 CPU 天。對 LLM Compiler FTD 進行指令微調以執行標誌調優任務的目標是在不需要執行編譯器數千次的情況下達到自動調優器效能的一部分。將程式碼從組合語言提升到更高層次的結構，可以執行額外的最佳化，例如直接整合到應用程式程式碼中的庫程式碼，或者將遺留程式碼移植到新架構。反編譯領域在將機器學習技術應用於從二進位制可執行檔案生成可讀和準確的程式碼方面取得了進展。在本研究中，研究團隊展示了 LLM Compiler FTD 如何透過微調進行反彙編，學習彙編程式碼和編譯器 IR 之間的關係。任務是學習 clang -xir - -o - -S 的逆向翻譯，如圖 5 所示。往返測試。使用 LLM 進行反彙編會導致正確性問題。提升的程式碼必須透過等價性檢查器進行驗證，這並不總是可行的，或者需要手動驗證正確性，或經過充分的測試用例以獲得信心。然而，可以透過往返測試找到正確性的下限。也就是說，透過將提升的 IR 重新編譯成彙編程式碼，如果彙編程式碼是相同的，則 IR 是正確的。這為使用 LLM 的結果提供了一條簡單途徑，並且是衡量反彙編模型效用的一種簡單方法。任務規範。研究團隊向模型提供彙編程式碼，並訓練它發出相應的反彙編 IR。這項任務的上下文長度設定為輸入彙編程式碼 8k 個 token 和輸出 IR8k 個 token。資料集。他們從之前任務中使用的資料集中派生出彙編程式碼和 IR 對。他們的微調資料集包含 470 萬個樣本，輸入 IR 在降低到 x86 彙編之前已經使用 - Oz 進行了最佳化。資料透過位元組對編碼進行標記化，使用與 Code Llama、Llama 和 Llama 2 相同的標記器。他們對所有四個訓練階段使用相同的訓練引數。他們使用的大部分訓練引數與 Code Llama 基礎模型相同，使用 AdamW 最佳化器，β1 和 β2 的值為 0.9 和 0.95。他們使用餘弦排程，預熱步驟為 1000 步，並將最終學習率設定為峰值學習率的 1/30。與 Code Llama 基礎模型相比，該團隊將單個序列的上下文長度從 4096 增加到 16384，但保持批次大小恆定為 400 萬個 token。為了適應更長的上下文，他們將學習率設定為 2e-5，並修改了 RoPE 位置嵌入的引數，其中他們將頻率重置為基本值 θ=10^6。這些設定與 Code Llama 基礎模型進行的長上下文訓練一致。該研究團隊評估 LLM Compiler 模型在標誌調優和反彙編任務、編譯器模擬、下一個 token 預測以及軟體工程任務上的表現。方法。他們評估 LLM Compiler FTD 在未見程式的最佳化標誌調優任務上的表現，並與 GPT-4 Turbo 和 Code Llama - Instruct 進行比較。他們對每個模型執行推理，從模型輸出中提取最佳化 pass 列表，然後他們使用這個 pass 列表來最佳化特定程式並記錄二進位制大小，基線是使用 -Oz 最佳化時程式的二進位制大小。對於 GPT-4 Turbo 和 Code Llama - Instruct，他們在提示後附加一個字尾，提供額外上下文以進一步描述問題和預期輸出格式。所有模型生成的 pass 列表都使用 PassListEval 進行驗證，如果驗證失敗則使用 -Oz 作為替代。為進一步驗證模型生成的 pass 列表的正確性，他們連結最終的程式二進位制檔案，並將其輸出與使用保守的 -O2 最佳化管道最佳化的基準輸出進行差分測試。資料集。研究團隊使用從 MiBench 基準套件提取的 2,398 個測試提示進行評估。為生成這些提示，他們取構成 24 個 MiBench 基準的所有 713 個翻譯單元，並從每個單元生成未最佳化的 IR，然後將它們格式化為提示。如果生成的提示超過 15k tokens，他們使用 llvm-extract 將代表該翻譯單元的 LLVM 模組分割成更小的模組，每個函式一個，這導致 1,985 個提示適合 15k token 上下文視窗，剩下 443 個翻譯單元不適合。在計算效能分數時，他們對 443 個被排除的翻譯單元使用 -Oz。表 10 總結了基準。結果。表 3 顯示了所有模型在標誌調優任務上的零樣本效能。只有 LLM Compiler FTD 模型比 -Oz 有所改進，13B 引數模型略優於較小的模型，在 61% 的情況下生成比 -Oz 更小的目標檔案。在某些情況下，模型生成的 pass 列表導致比 -Oz 更大的目標檔案大小。例如，LLM Compiler FTD 13B 在 12% 的情況下有退化。這些退化可以透過簡單地編譯程式兩次來避免：一次使用模型生成的 pass 列表，一次使用 -Oz，然後選擇產生最佳結果的 pass 列表。透過消除相對於 -Oz 的退化，這些 -Oz 備份分數將 LLM Compiler FTD 13B 相對於 -Oz 的總體改進提高到 5.26%，並使 Code Llama - Instruct 和 GPT-4 Turbo 相對於 -Oz 有適度的改進。圖 6 顯示了每個模型在各個基準上的效能細分。 二進位制大小準確性。雖然模型生成的二進位制大小預測對實際編譯沒有影響，但研究團隊可以評估模型在預測最佳化前後的二進位制大小方面的效能，以瞭解每個模型對最佳化的理解程度。圖 7 顯示了結果。LLM Compiler FTD 的二進位制大小預測與實際情況相關性良好，7B 引數模型對未最佳化和最佳化的二進位制大小分別達到了 0.083 和 0.225 的 MAPE 值。13B 引數模型的 MAPE 值相似，分別為 0.082 和 0.225。Code Llama - Instruct 和 GPT-4 Turbo 的二進位制大小預測與實際情況幾乎沒有相關性。研究人員注意到，LLM Compiler FTD 對最佳化程式碼的錯誤略高於未最佳化程式碼。特別是 LLM Compiler FTD 偶爾有高估最佳化效果的趨勢，導致預測的二進位制大小低於實際情況。消融研究。表 4 對模型在 500 個提示的小型保留驗證集上的效能進行了消融研究，這些提示來自與他們訓練資料相同的分佈 (但未在訓練中使用)。他們在圖 1 所示訓練管道的每個階段進行標誌調優訓練，以比較效能。如圖所示，反彙編訓練導致效能從平均 5.15% 略微下降到 5.12%(相對於 -Oz 的改進)。他們還展示了用於生成第 2 節所述訓練資料的自動調優器的效能。LLM Compiler FTD 達到了自動調優器 77% 的效能。方法。研究團隊評估 LLM 生成的程式碼在將彙編程式碼反彙編到 LLVM-IR 時的功能正確性。他們評估 LLM Compiler FTD 並與 Code Llama - Instruct 和 GPT-4 Turbo 進行比較，發現需要額外的提示字尾才能從這些模型中提取最佳效能。字尾提供了關於任務和預期輸出格式的額外上下文。為評估模型的效能，他們將模型生成的反彙編 IR 往返降級回彙編。這使我們能夠透過比較原始彙編與往返結果的 BLEU 分數來評估反彙編的準確性。從彙編到 IR 的無損完美反彙編將有 1.0 的往返 BLEU 分數 (精確匹配)。資料集。他們使用從 MiBench 基準套件提取的 2,015 個測試提示進行評估，取用於上述標誌調優評估的 2,398 個翻譯單元，生成反彙編提示。然後他們根據最大 8k token 長度過濾提示，允許 8k tokens 用於模型輸出，剩下 2,015 個。表 11 總結了基準。LLM Compiler FTD 7B 的往返成功率略高於 LLM Compiler FTD 13B，但 LLM Compiler FTD 13B 具有最高的往返彙編準確性 (往返 BLEU) 和最頻繁產生完美反彙編 (往返精確匹配)。Code Llama - Instruct 和 GPT-4 Turbo 在生成語法正確的 LLVM-IR 方面存在困難。圖 8 顯示了所有模型的往返 BLEU 分數分佈。消融研究。表 6 對模型在 500 個提示的小型保留驗證集上的效能進行了消融研究，這些提示取自之前使用的 MiBench 資料集。他們在圖 1 所示訓練管道的每個階段進行反彙編訓練，以比較效能。往返率在透過整個訓練資料堆疊時最高，並隨每個訓練階段持續下降，儘管往返 BLEU 在每個階段變化不大。方法。該研究團隊在下一個 token 預測和編譯器模擬兩個基礎模型任務上對 LLM Compiler 模型進行消融研究。他們在訓練管道的每個階段進行這種評估，以瞭解為每個連續任務訓練如何影響效能。對於下一個 token 預測，他們在所有最佳化級別的 LLVM-IR 和彙編程式碼的小樣本上計算困惑度。他們使用兩個指標評估編譯器模擬：生成的 IR 或彙編程式碼是否編譯，以及生成的 IR 或彙編程式碼是否與編譯器產生的完全匹配。資料集。對於下一個 token 預測，他們使用從與我們訓練資料相同分佈但未用於訓練的小型保留驗證資料集。他們使用混合的最佳化級別，包括未最佳化程式碼、用 -Oz 最佳化的程式碼和隨機生成的 pass 列表。對於編譯器模擬，他們使用從 MiBench 生成的 500 個提示進行評估，這些提示使用第 2.2 節描述的方式隨機生成的 pass 列表。結果。表 7 顯示了 LLM Compiler FTD 在所有訓練階段在兩個基礎模型訓練任務 (下一個 token 預測和編譯器模擬) 上的效能。下一個 token 預測效能在 Code Llama 之後急劇上升，後者幾乎沒有見過 IR 和彙編，並在隨後的每個微調階段略有下降。對於編譯器模擬，Code Llama 基礎模型和預訓練模型表現不佳，因為它們沒有在這個任務上訓練過。在編譯器模擬訓練之後直接達到最高效能，其中 LLM Compiler FTD 13B 生成的 95.6% 的 IR 和彙編可以編譯，20% 與編譯器完全匹配。在進行標誌調優和反彙編微調後，效能下降。方法。雖然 LLM Compiler FTD 的目的是為程式碼最佳化提供基礎模型，但它建立在為軟體工程任務訓練的基礎 Code Llama 模型之上。為評估 LLM Compiler FTD 的額外訓練如何影響程式碼生成的效能，他們使用與 Code Llama 相同的基準套件，評估 LLM 從自然語言提示生成 Python 程式碼的能力，如「編寫一個函式，找出可以從給定的對集合形成的最長鏈。」資料集。他們使用 HumanEval 和 MBPP 基準，與 Code Llama 相同。結果。表 8 顯示了從 Code Llama 基礎模型開始的所有模型訓練階段和模型大小的貪婪解碼效能 (pass@1)。它還顯示了模型在 pass@10 和 pass@100 上的分數，這些分數是用 p=0.95 和 temperature=0.6 生成的。每個以編譯器為中心的訓練階段都導致 Python 程式設計能力略有退化。在 HumanEval 和 MBPP 上，LLM Compiler 的 pass@1 效能最多下降 18% 和 5%，LLM Compiler FTD 在額外的標誌調優和反彙編微調後最多下降 29% 和 22%。所有模型在這兩個任務上仍然優於 Llama 2。Meta 研究團隊已經展示了 LLM Compiler 在編譯器最佳化任務上表現良好，並且相比先前的工作，對編譯器表示和彙編程式碼的理解有所改進，但仍存在一些侷限性。主要限制是輸入的有限序列長度 (上下文視窗)。LLM Compiler 支援 16k tokens 的上下文視窗，但程式程式碼可能遠遠超過這個長度。例如，當格式化為標誌調優提示時，67% 的 MiBench 翻譯單元超過了這個上下文視窗，如表 10 所示。為了緩解這一問題，他們將較大的翻譯單元拆分為單獨的函式，儘管這限制了可以執行的最佳化範圍，而且仍有 18% 的拆分翻譯單元對模型來說太大，無法作為輸入接受。研究人員正在採用不斷增加的上下文視窗，但有限的上下文視窗仍然是 LLM 的一個普遍問題。第二個限制，也是所有 LLM 的共同問題，是模型輸出的準確性。建議 LLM Compiler 的使用者使用特定於編譯器的評估基準來評估他們的模型。鑑於編譯器並非無 bug，任何建議的編譯器最佳化都必須經過嚴格測試。當模型反編譯彙編程式碼時，其準確性應透過往返、人工檢查或單元測試來確認。對於某些應用，LLM 生成可以被限制在正規表示式內，或與自動驗證相結合以確保正確性。https://x.com/AIatMeta/status/1806361623831171318
https://ai.meta.com/research/publications/meta-large-language-model-compiler-foundation-models-of-compiler-optimization/?utm_source=twitter&utm_medium=organic_social&utm_content=link&utm_campaign=fair