如何訓練最強代碼大模型?北大aiXcoder-7B貢獻前沿實踐
本文的通訊作者是北京大學計算機學院長聘教授李戈。
本文一作是 aiXcoder 蔣思源和北大李戈教授課題組博士生李佳,團隊重點關注融合深度學習與軟件工程的代碼建模方法。
如何訓練一個代碼大模型?這一過程看似簡單:獲取代碼數據、清洗數據,最終啟動訓練。如今,開源代碼數據集層出不窮;數據清洗工具也已成熟,包括開源的許可證識別工具、MinHash 算法、PII 識別模型等;而在分布式訓練方面,像 Megatron-LM、DeepSpeed 等框架也大大降低了技術門檻。看似我們只差計算資源,就能訓練出一個強大的代碼大模型。
然而,訓練模型的初衷,應該始終從實際開發場景出發。作為開發者,我們不僅需要了解定義的各種 API 接口,還需要從入口函數模擬程序的執行過程,追蹤到每一行修改的代碼。在復雜的項目中,任何小小的變動都可能影響整個系統的運轉。
但現有的代碼大模型并未充分考慮到軟件開發的具體場景,它們往往將最終版本的代碼簡單地視作自然語言文本,試圖通過復制自然語言處理的成功經驗來處理代碼。這種方法忽略了代碼的結構性和復雜的上下文關系,導致模型在實際開發中表現不佳。
北京大學 aiXcoder 團隊一直致力于探索如何將深度學習與軟件開發深度融合,推動軟件開發的自動化。2024 年 4 月,aiXcoder 開源了自研代碼大模型 aiXcoder-7B,成為這一領域的一次重要嘗試,旨在將代碼的抽象語法樹(AST)結構與大規模預訓練結合,以期提升模型對代碼結構和上下文的理解能力。
近期,該篇論文被軟件工程領域國際頂級會議 ICSE 2025 收錄,將于 4 月 27 日 - 5 月 3 日赴加拿大渥太華參會分享研究成果。
此次論文錄用不僅是對 aiXcoder 7B 代碼大模型技術前瞻性和應用創新性的高度認可,更標志著該模型繼成功落地企業并獲各行業客戶廣泛認可后,再次于學術界獲得權威肯定,充分彰顯了 aiXcoder 在推動軟件工程發展中的前瞻性引領作用。
- 論文地址:https://arxiv.org/pdf/2410.13187
- 開源項目地址:https://github.com/aixcoder-plugin/aiXcoder-7B
代碼數據,異于自然語言
相較于自然語言文本,程序是現實世界解決方案在計算機系統中的映射。因此,程序源代碼呈現出很多獨特的性質,例如:強結構性、可執行性等等。有效地表示和建模這些特性,對于代碼生成等任務來說至關重要。
如上三行代碼能夠嚴格解析為抽象語法樹格式
代碼天然能被解析為抽象語法樹,其語法規則嚴格組織了代碼語句之間的關系。在語法規則之上,也有很多方式描述代碼之間的流轉關系,例如控制流圖、調用流圖等等。顧名思義,控制流圖會展示整個代碼控制與條件關系,什么樣的條件下哪個分支代碼會運行。調用流圖則展示的是代碼之間的調用關系,實現一個功能時在什么樣的地方調用什么樣的代碼模塊是能展示出來的。
控制流圖示例,代碼執行條件與順序會解析成流程圖。
調用流圖示例,main 函數調用 calculate 函數計算兩個數之和,calculate 函數調用另外兩個函數 getFirst 和 getSecond 獲取參與計算的兩個加數。
程序語言與自然語言之間存在顯著差異。盡管大模型通過大規模自回歸訓練任務在通用知識學習上取得了巨大成功,但這并不意味著可以簡單地將代碼數據視為「自然語言」,并將其拉長為一維 Token 序列進行自回歸訓練,就能復制自然語言處理的成功。
事實上,當使用自回歸模型或「Fill in the middle」任務訓練基礎模型時,會發現實際在代碼補全任務中,模型生成的結果往往與人類程序員的編程方式不符,我們還需要更符合代碼的預訓練方法。
aiXcoder-7B:創新在 LLM 上引入代碼特性
正因為當前代碼大模型很少將代碼特性引入到 LLM 的訓練過程中,代碼大模型在企業真實項目中表現得不盡人意,所以我們創新將一些傳統軟件工程方法引入到大規模預訓練中,希望能生成更符合真實場景的代碼內容。
為此,aiXcoder-7B 主要從以下幾個方面優化預訓練:
- 數據預處理:軟工工具保證代碼數據語法正確且不存在嚴重 Bug
- 結構化 FIM:按照語法結構組織預訓練任務
- 多文件排序:保證單項目內,文件排序既考慮內容相似,又考慮調用關系
數據預處理
aiXcoder 核心數據集主要用于強化代碼大模型在以上編程語言上的效果,其經過大量的過濾與篩選過程。相比于其它代碼大模型,aiXcoder-7B 預訓練數據既采用常規的數據處理,例如數據去重、自動生成代碼去除、通過 Star 量、正則等規則去除低質量代碼、敏感信息等,同時借助軟件工程方法進行更精細的數據處理。
具體而言,aiXcoder-7B 預訓練數據采用語法分析和靜態分析兩大類工具預處理數據。對于語法分析,重點解析五十種主流語言的語法結構,并排除存在語法錯誤、簡單 Bug 、大面積被注釋掉的代碼等。
語法分析能天然解析并處理明顯不合理的代碼
對于靜態分析,則側重解析十余種最主流編程語言的嚴重錯誤,即當出現這一些類型錯誤時,代碼大概率在執行過程中會出現比較大的問題。具體而言,掃描并定位影響代碼可靠性和可維護性的 161 種 Bug,影響代碼安全性的 197 種安全漏洞。
靜態分析能檢測出很多更深層缺陷與漏洞的代碼。
結合軟件工程分析方法以及過濾規則,能夠將存在明顯問題的代碼刪除掉,明顯提升整體代碼質量。
結構化 FIM
在實際開發過程中,代碼具有類、方法、條件代碼塊、循環代碼塊等眾多結構。研究團隊期待讓代碼大模型天然能學會這樣的結構,而不是放任代碼大模型向下一直生成,或者從字符層面上截取一個片段,期待補全該字符片段。
為此,團隊結合語法分析方法,將代碼解析為抽象語法樹,并基于語法樹的結構構建訓練任務。具體而言,代碼文件中的每個位置都對應著抽象語法樹中的某個節點。在訓練過程中,團隊挖掉該節點的子節點,或者挖掉該節點所在父節點剩余的部分,然后針對被挖掉的代碼塊做一個先驗約束:挖掉的代碼塊橫跨一個或少數幾個完整的代碼結構。將這部分完整代碼結構用來計算損失訓練模型,就能一定程度上讓代碼模型理解部分語法結構。
更形象地解釋,常規的 Fill in the middle 會構造很多不合法的代碼片段,例如下圖「or i in range (2」,常規的做法只是從字符上隨機取一個片段。但論文研究團隊提出的 Structured Fill-In-the-Middle (SFIM) 會隨機先選定一個語法節點「IF」,并在 IF 節點向下取了「Compare」代碼片段「i % 5 == 0:」
最終團隊在預訓練中根據 SFIM 構建整體訓練損失計算,以此更好地學習代碼的語法結構信息。
多文件排序
當前主流的代碼開源數據集,例如 TheStackV1 、 TheStackV2 或者 The Pile 中代碼部分,都是根據單個語言,甚至單一后綴名組織數據,致使整個訓練樣本的構造局限在單語言文件中。而此次研究團隊構建的訓練數據以項目為單位,保留與處理多種編程語言的代碼文件,確保訓練數據中編程語言的分布與真實開發一致。
此時有一個重要的問題:項目內不同文件該如何排序?
為了提升模型對項目內多代碼文件關系的充分建模能力,并在推理過程中更高效地抽取有用的上下文信息,研究團隊通過相似性關系和依賴性關系對代碼文件排序:相似性關系即模型在預訓練中能學會仿寫相似的代碼;依賴性關系即模型在預訓練中能學會 API 調用或者函數調用的關系。
預訓練中,項目內文件排序算法
如 Algorithm 1 所示,本論文給出了一個項目內文件排序偽代碼。簡單理解,以 0.3 的概率采用文件內容相似排序,即通過 KMeans 聚類算法將文件聚成不同的簇,并且同一個簇排列在一起;此外,以 0.3 的概率進行路徑相似排序,把同一目錄下的文件,或者被測代碼與測試代碼等路徑相關的文件能排列在一起;最后還以 0.3 的概率構建函數調用流圖,并根據圖的葉節點一路向根節點建立程序依賴路徑,將路徑上的代碼文件排列在一起。
aiXcoder 7B 獨特的效果優勢
借助軟件工程方法,研究團隊通過更符合代碼大模型的預訓練方法,提升了其在代碼數據上的理解與生成能力。例如論文表 5 中的 Fill-in-the-middle 評測集顯示,經過高質量代碼數據的 SFIM 任務訓練,不同語言的代碼補全能力有明顯的提升。
為了進一步測評 aixcoder-7B 在多種情況下的代碼補全能力,團隊從方法簽名、方法體、方法局部、條件塊、循環塊、異常捕捉塊等維度評估了模型在代碼補全上的效果,如論文圖 4 所示。對比 DeepSeekcoder-6.7B,aixcoder-7B 大部分的補全位置都擁有更好的效果。
此外,因為預訓練任務充分考慮了代碼的語法結構,模型在推理過程中對代碼的上下文結構展現出更出色的感知能力,能夠準確判斷需要補全完整的語法結構,并傾向生成更短的代碼片段。如論文表 6 所示,模型生成的 Token 數與 GroundTruth Token 數的比值,aiXcoder -7B 更小,表明 SFIM 預訓練任務有效指導了模型更好學會如何終止預測。
對于代碼補全任務,另外一個比較重要的是跨文件上下文的理解能力。aiXcoder -7b 在預訓練中以項目為單位對項目內的代碼文件進行排序,獲得了更好的文件間建模能力。如表 4 所示,aiXcoder -7b 在 CrossCodeEval 評測集上擁有更好的效果,表明其利用多文件的上下文信息,補全當前代碼文件能力更有優勢。
后續改進方向
在真實軟件開發場景中,還有很多能力是大模型未曾學習到的,重中之重即代碼上下文。
實際代碼補全往往需要基于不同類型的上下文(如:當前文件的上文、跨文件上下文、相似代碼段),去預測后續的代碼。這種復雜的上下文形式與基礎模型預訓練時的上下文形式不一致,從而限制了基礎模型在實際應用時的代碼補全準確率。
為解決這個問題,研究團隊在 aiXcoder 7B 上做了更多的對齊訓練實驗。該對齊訓練有效地將模型對齊到真實軟件開發場景中的上下文形式,顯著地提升了模型在多種語言上的代碼補全準確率。例如,在四種語言(Python、Java、C++和Go)的多行補全上,相較于aiXcoder-7B,經過優化的新模型在Exact Match(完全匹配)指標上平均取得了 13 個點的絕對提升。
當前,充分利用數十年積累的軟件工程經驗,將代碼大模型真正應用于軟件開發的實際場景中,仍然是一項艱巨而復雜的任務。然而,隨著不斷深入的研究,代碼大模型已經讓「軟件開發自動化」這一宏偉目標變得愈加觸手可及。