KwaiCoder-23BA4-v1:以 1/30 的成本訓練全尺寸 SOTA 代碼續寫大模型 原創
KwaiCoder-23BA4-v1 是快手 Kwaipilot 團隊最新開源的自研代碼續寫大模型,模型的訓練依托于 Kwaipilot 團隊提出的一種高效的模型訓練方案。通過結合模型剪枝、知識蒸餾、細粒度合并等技術,相比傳統方法以 1/30 的成本完成了 23B 寬 MoE 架構代碼續寫模型的訓練,并在多個代碼領域評測集上實現了新的 SOTA。
開源地址:??https://huggingface.co/Kwaipilot/KwaiCoder-23B-A4B-v1??
我們的路徑:
- 高質量數據:我們強調篩選和整理高質量代碼、文本與合成數據,而非單純擴大數據集規模,從而確保數據的相關性和一致性;
- 一種經濟且高效的模型訓練路線:通過結合模型知識蒸餾與細粒度合并,完成模型的知識壓縮(scale down)與模型參數擴充(scale up)。在 scale down 階段,模型參數減少,模型知識密度得到了提升,再基于此完成模型的剪裁與合并,使得模型以一種極高的知識密度完成了參數量的擴充,從而實現了性能的跨越。
結果表明,該方法不僅超越了基線模型的表現,還在多個基準上達到了 SOTA 指標。我們在整個訓練路線中進行了 350 多次大小實驗,本報告不僅詳細介紹了我們的方法、也會將實驗中收獲的認知分享給大家,在文章中加粗體現。
模型性能
KwaiCoder-23BA4-v1 在 HumanEval 和 HumanEval+測試集上的 Pass@1 指標上分別達到 82.9%和 76.2%,超過了目前最好的 Base 模型(OpenCoder-8B),達到了 SOTA 水平。
我們也在 BigCodeBench 上進行了評測,BigCodeBench 是一個用于評估大型語言模型在解決實際和具有挑戰性的編程任務上的性能的基準測試,在 BigCodeBench-Complete 全集和 Hard 子集上,KwaiCoder-23BA4-v1 以 4B 的激活參數量僅次于 Qwen2.5 Coder 32B base。
在多語言能力上(Multipl-e), 也超過了目前最好的 Base 模型(OpenCoder-8B,62.8),達到了 SOTA 水平。
在代碼續寫任務上,我們測評了 Fill-in-Middle 榜單,超過了目前最好的 Base 模型 (Qwen2.5 Coder 32B-Base, 88.3),達到了 SOTA 水平。
數據構造
我們構造了 3T 的高質量預訓練數據集,包括代碼、數學與知識類文本語料。對于源碼數據,我們依照研發場景的實際使用情況對編程語言分布和過濾策略進行了優化。為了提升語料中的知識濃度,我們使用了基于模型的過濾方法,并追加了更多高質量合成數據。多階段訓練中的詳細語料分布如下圖所示(左預訓練,右退火)
面向代碼續寫補全能力,我們使用了在線 Fill-in-Middle (FIM) 策略,并對切分方式進行了優化,在保障樣本完整和多樣性的同時避免截斷樣本帶來的幻覺問題。我們發現對切分方式的精細調整可以對補全能力帶來較大的提升,使用經過調整后的格式僅在幾百步訓練中就達到了 Fill-in-middle 榜單的 SOTA 的水平(88->93)。
對于開放源碼,我們利用倉庫級別的程序分析技術以及細粒度策略規則進行了精細清洗。基于啟發式策略過濾數據文件、環境文件、損壞或統計信號異常的文件;根據不同編程語言的特點來過濾掉邏輯不完整、格式異常的代碼文件;對隱私子串與評測集污染進行過濾;利用程序分析構建倉庫級別樣本,并分別進行文件級與倉庫級的精確和模糊去重,最大程度減小數據中的冗余并保持多樣性。
在基于模型的處理中,我們構建了多維度的數據質量信號,并使用多級分類模型從開放源碼中提取高質量子集;對于知識類語料,我們使用簡單模型在開放語料中進行聚類挖掘,并進行多級過濾。我們的實驗表明,對數據進行精細的多階段過濾,能夠顯著提升模型的效果。但是這也使得我們的訓練數據缺乏,在模型的訓練階段,我們觀察到模型下游任務的上漲趨勢并沒有放緩。我們嘗試混入一些沒有經過精細過濾的數據,但是這部分數據并不能讓模型獲得更好的性能,我們不得不提前停止了模型的訓練。
我們利用參數量更大、能力更強的教師模型進行了高質量合成,以實現對數理領域的定向優化。我們通過主題增強和結構優化,生產內容覆蓋更加完善、由淺入深且結構清晰的教科書語料。以代碼相關問答數據為種子,我們生產了包含短推理鏈的合成源碼片段,并結合測試用例驗證以消除模型幻覺。
模型訓練
我們認為,目前大部分的開源模型在冷啟動階段 (from scratch) 使用了相似且巨量的數據、這一階段的數據由于數據量巨大往往不會采用精細的數據過濾方法,數據質量一般,且資源消耗極大。
因此許多模型訓練廠商采用基于 base 模型直接進行 cpt 的方式完成模型訓練的需求。但是這種方案無法按需更改模型的結構、tokenzier、Attention 實現,也難以根據自身的模型集群狀況 (卡的型號,顯存帶寬的大小等) 制定最合理的參數規模。
于是為了完成我們對模型結構的定制化需求,權衡模型的訓練成本,我們決定使用 Pruning(模型裁剪) + 知識蒸餾(Knowledge Distillation)+ 細粒度合并(Granular Upcyling) 的技術路線。
模型裁剪
我們首先對模型的 MLP、LayerNorm、Attention Head、Layer 等進行重要性估計,并按照估計的結果對模型參數進行裁剪。我們對許多優秀的業界開源模型進行了實驗,我們發現大部分開源模型的知識密度是不飽和的,這一點尤其體現在模型中間靠后的層中。
我們實驗發現,雖然在前期寬度剪枝的結果會優于深度剪枝,但是在后續因為我們要對模型的寬度再次切分,因此在前期進行過度的寬度剪枝會一定程度影響模型的最終性能。
?
模型蒸餾
再將裁剪前的模型作為教師模型,裁剪后的模型作為學生模型進行在線知識蒸餾。對模型的中間結果與 logits 進行蒸餾,以進一步縮小教師模型與學生模型的差距。
在我們最后的結果中,學生模型以接近一半的參數量在 HumanEval、Fill-in-Middle、BigcodeBench、GSM8K、MBPP 等榜單上達成了教師模型 95%以上的性能。我們的實驗結果表明,在蒸餾之前對 Teacher Model 使用與蒸餾相同的數據進行微調,可以大幅提升蒸餾的效率。
?
細粒度合并
我們進一步將學生模型按照寬度進行更細粒度切分,以完成細粒度的專家構造與寬體 MoE 的合并,在這個過程中,我們對模型參數進行了縮放,以確保合并之后的 MoE 模型是無損的。
- 模型切分
- 參數縮放
- 門控初始化
我們也做了許多充分的消融實驗 (>100B token) ,我們的實驗發現:這種合并方式不僅能夠使得 CPT 的初始階段 loss 更加穩定,并且相比傳統 UpCycling 或是直接 CPT 有更高的能力上限。
?
CPT
在最后的 CPT 階段,我們使用了三階段的訓練方法,同時我們將數據配比與學習率公布出來,以便于后續模型訓練工作者完成模型的二次開發。
我們將最后的模型在快手內部 kwaipilot 產品上上線,我們觀察到采納率指標相比 DeepSeek-Coder-V2-Lite-Base 提升了 4-5pp 左右。
總結與展望
KwaiCoder-23BA4-v1 模型的發布是我們在代碼生成和續寫領域探索的重要里程碑,作為一種經濟且高效獲取自定義模型的方案,突破了這條路線的技術瓶頸。在未來的研究中,我們計劃從以下幾個方面進一步推進:
- 對模型剪枝和知識壓縮方案的進一步探索;
- 利用技術路線中可自定義模型結構的優勢,探索更高效的訓練、推理算法方案;
- 更廣泛的開源和社區合作,我們將繼續推進模型和方案的開放性,并與學術界和工業界合作,共同構建更強大的代碼生成系統。
引用
[1] Muralidharan S, Sreenivas S T, Joshi R B, et al. Compact language models via pruning and knowledge distillation[C]//The Thirty-eighth Annual Conference on Neural Information Processing Systems. 2024.
[2] He E, Khattar A, Prenger R, et al. Upcycling large language models into mixture of experts[J]. arXiv preprint arXiv:2410.07524, 2024.
[3] Zhu Q, Guo D, Shao Z, Yang D, Wang P, Xu R, Wu Y, Li Y, Gao H, Ma S, Zeng W. DeepSeek-Coder-V2: Breaking the Barrier of Closed-Source Models in Code Intelligence. arXiv preprint arXiv:2406.11931. 2024 Jun 17.
[4] Huang S, Cheng T, Liu JK, Hao J, Song L, Xu Y, Yang J, Liu JH, Zhang C, Chai L, Yuan R. Opencoder: The open cookbook for top-tier code large language models. arXiv preprint arXiv:2411.04905. 2024 Nov 7.
