KwaiCoder-23BA4-v1 是快手 Kwaipilot 團隊最新開源的自研代碼續(xù)寫大模型，模型的訓(xùn)練依托于 Kwaipilot 團隊提出的一種高效的模型訓(xùn)練方案。通過結(jié)合模型剪枝、知識蒸餾、細(xì)粒度合并等技術(shù)，相比傳統(tǒng)方法以 1/30 的成本完成了 23B 寬 MoE 架構(gòu)代碼續(xù)寫模型的訓(xùn)練，并在多個代碼領(lǐng)域評測集上實現(xiàn)了新的 SOTA。

開源地址：??https://huggingface.co/Kwaipilot/KwaiCoder-23B-A4B-v1??

我們的路徑：

1. 高質(zhì)量數(shù)據(jù)：我們強調(diào)篩選和整理高質(zhì)量代碼、文本與合成數(shù)據(jù)，而非單純擴大數(shù)據(jù)集規(guī)模，從而確保數(shù)據(jù)的相關(guān)性和一致性；
2. 一種經(jīng)濟且高效的模型訓(xùn)練路線：通過結(jié)合模型知識蒸餾與細(xì)粒度合并，完成模型的知識壓縮(scale down)與模型參數(shù)擴充(scale up)。在 scale down 階段，模型參數(shù)減少，模型知識密度得到了提升，再基于此完成模型的剪裁與合并，使得模型以一種極高的知識密度完成了參數(shù)量的擴充，從而實現(xiàn)了性能的跨越。

結(jié)果表明，該方法不僅超越了基線模型的表現(xiàn)，還在多個基準(zhǔn)上達(dá)到了 SOTA 指標(biāo)。我們在整個訓(xùn)練路線中進(jìn)行了 350 多次大小實驗，本報告不僅詳細(xì)介紹了我們的方法、也會將實驗中收獲的認(rèn)知分享給大家，在文章中加粗體現(xiàn)。

模型性能

KwaiCoder-23BA4-v1 在 HumanEval 和 HumanEval+測試集上的 Pass@1 指標(biāo)上分別達(dá)到 82.9%和 76.2%，超過了目前最好的 Base 模型（OpenCoder-8B），達(dá)到了 SOTA 水平。

我們也在 BigCodeBench 上進(jìn)行了評測，BigCodeBench 是一個用于評估大型語言模型在解決實際和具有挑戰(zhàn)性的編程任務(wù)上的性能的基準(zhǔn)測試，在 BigCodeBench-Complete 全集和 Hard 子集上，KwaiCoder-23BA4-v1 以 4B 的激活參數(shù)量僅次于 Qwen2.5 Coder 32B base。

在多語言能力上(Multipl-e), 也超過了目前最好的 Base 模型（OpenCoder-8B，62.8），達(dá)到了 SOTA 水平。

在代碼續(xù)寫任務(wù)上，我們測評了 Fill-in-Middle 榜單，超過了目前最好的 Base 模型 (Qwen2.5 Coder 32B-Base, 88.3)，達(dá)到了 SOTA 水平。

數(shù)據(jù)構(gòu)造

我們構(gòu)造了 3T 的高質(zhì)量預(yù)訓(xùn)練數(shù)據(jù)集，包括代碼、數(shù)學(xué)與知識類文本語料。對于源碼數(shù)據(jù)，我們依照研發(fā)場景的實際使用情況對編程語言分布和過濾策略進(jìn)行了優(yōu)化。為了提升語料中的知識濃度，我們使用了基于模型的過濾方法，并追加了更多高質(zhì)量合成數(shù)據(jù)。多階段訓(xùn)練中的詳細(xì)語料分布如下圖所示（左預(yù)訓(xùn)練，右退火）

面向代碼續(xù)寫補全能力，我們使用了在線 Fill-in-Middle (FIM) 策略，并對切分方式進(jìn)行了優(yōu)化，在保障樣本完整和多樣性的同時避免截斷樣本帶來的幻覺問題。我們發(fā)現(xiàn)對切分方式的精細(xì)調(diào)整可以對補全能力帶來較大的提升，使用經(jīng)過調(diào)整后的格式僅在幾百步訓(xùn)練中就達(dá)到了 Fill-in-middle 榜單的 SOTA 的水平(88->93)。

對于開放源碼，我們利用倉庫級別的程序分析技術(shù)以及細(xì)粒度策略規(guī)則進(jìn)行了精細(xì)清洗?；趩l(fā)式策略過濾數(shù)據(jù)文件、環(huán)境文件、損壞或統(tǒng)計信號異常的文件；根據(jù)不同編程語言的特點來過濾掉邏輯不完整、格式異常的代碼文件；對隱私子串與評測集污染進(jìn)行過濾；利用程序分析構(gòu)建倉庫級別樣本，并分別進(jìn)行文件級與倉庫級的精確和模糊去重，最大程度減小數(shù)據(jù)中的冗余并保持多樣性。

在基于模型的處理中，我們構(gòu)建了多維度的數(shù)據(jù)質(zhì)量信號，并使用多級分類模型從開放源碼中提取高質(zhì)量子集；對于知識類語料，我們使用簡單模型在開放語料中進(jìn)行聚類挖掘，并進(jìn)行多級過濾。我們的實驗表明，對數(shù)據(jù)進(jìn)行精細(xì)的多階段過濾，能夠顯著提升模型的效果。但是這也使得我們的訓(xùn)練數(shù)據(jù)缺乏，在模型的訓(xùn)練階段，我們觀察到模型下游任務(wù)的上漲趨勢并沒有放緩。我們嘗試混入一些沒有經(jīng)過精細(xì)過濾的數(shù)據(jù)，但是這部分?jǐn)?shù)據(jù)并不能讓模型獲得更好的性能，我們不得不提前停止了模型的訓(xùn)練。

我們利用參數(shù)量更大、能力更強的教師模型進(jìn)行了高質(zhì)量合成，以實現(xiàn)對數(shù)理領(lǐng)域的定向優(yōu)化。我們通過主題增強和結(jié)構(gòu)優(yōu)化，生產(chǎn)內(nèi)容覆蓋更加完善、由淺入深且結(jié)構(gòu)清晰的教科書語料。以代碼相關(guān)問答數(shù)據(jù)為種子，我們生產(chǎn)了包含短推理鏈的合成源碼片段，并結(jié)合測試用例驗證以消除模型幻覺。

模型訓(xùn)練

我們認(rèn)為，目前大部分的開源模型在冷啟動階段 (from scratch) 使用了相似且巨量的數(shù)據(jù)、這一階段的數(shù)據(jù)由于數(shù)據(jù)量巨大往往不會采用精細(xì)的數(shù)據(jù)過濾方法，數(shù)據(jù)質(zhì)量一般，且資源消耗極大。

因此許多模型訓(xùn)練廠商采用基于 base 模型直接進(jìn)行 cpt 的方式完成模型訓(xùn)練的需求。但是這種方案無法按需更改模型的結(jié)構(gòu)、tokenzier、Attention 實現(xiàn)，也難以根據(jù)自身的模型集群狀況 (卡的型號，顯存帶寬的大小等) 制定最合理的參數(shù)規(guī)模。

于是為了完成我們對模型結(jié)構(gòu)的定制化需求，權(quán)衡模型的訓(xùn)練成本，我們決定使用 Pruning(模型裁剪) + 知識蒸餾(Knowledge Distillation）+ 細(xì)粒度合并(Granular Upcyling) 的技術(shù)路線。

模型裁剪

我們首先對模型的 MLP、LayerNorm、Attention Head、Layer 等進(jìn)行重要性估計，并按照估計的結(jié)果對模型參數(shù)進(jìn)行裁剪。我們對許多優(yōu)秀的業(yè)界開源模型進(jìn)行了實驗，我們發(fā)現(xiàn)大部分開源模型的知識密度是不飽和的，這一點尤其體現(xiàn)在模型中間靠后的層中。

我們實驗發(fā)現(xiàn)，雖然在前期寬度剪枝的結(jié)果會優(yōu)于深度剪枝，但是在后續(xù)因為我們要對模型的寬度再次切分，因此在前期進(jìn)行過度的寬度剪枝會一定程度影響模型的最終性能。

?

模型蒸餾

再將裁剪前的模型作為教師模型，裁剪后的模型作為學(xué)生模型進(jìn)行在線知識蒸餾。對模型的中間結(jié)果與 logits 進(jìn)行蒸餾，以進(jìn)一步縮小教師模型與學(xué)生模型的差距。

在我們最后的結(jié)果中，學(xué)生模型以接近一半的參數(shù)量在 HumanEval、Fill-in-Middle、BigcodeBench、GSM8K、MBPP 等榜單上達(dá)成了教師模型 95%以上的性能。我們的實驗結(jié)果表明，在蒸餾之前對 Teacher Model 使用與蒸餾相同的數(shù)據(jù)進(jìn)行微調(diào)，可以大幅提升蒸餾的效率。

?

細(xì)粒度合并

我們進(jìn)一步將學(xué)生模型按照寬度進(jìn)行更細(xì)粒度切分，以完成細(xì)粒度的專家構(gòu)造與寬體 MoE 的合并，在這個過程中，我們對模型參數(shù)進(jìn)行了縮放，以確保合并之后的 MoE 模型是無損的。

1. 模型切分

2. 參數(shù)縮放

3. 門控初始化

我們也做了許多充分的消融實驗 (>100B token) ，我們的實驗發(fā)現(xiàn)：這種合并方式不僅能夠使得 CPT 的初始階段 loss 更加穩(wěn)定，并且相比傳統(tǒng) UpCycling 或是直接 CPT 有更高的能力上限。

?

CPT

在最后的 CPT 階段，我們使用了三階段的訓(xùn)練方法，同時我們將數(shù)據(jù)配比與學(xué)習(xí)率公布出來，以便于后續(xù)模型訓(xùn)練工作者完成模型的二次開發(fā)。

我們將最后的模型在快手內(nèi)部 kwaipilot 產(chǎn)品上上線，我們觀察到采納率指標(biāo)相比 DeepSeek-Coder-V2-Lite-Base 提升了 4-5pp 左右。

總結(jié)與展望

KwaiCoder-23BA4-v1 模型的發(fā)布是我們在代碼生成和續(xù)寫領(lǐng)域探索的重要里程碑，作為一種經(jīng)濟且高效獲取自定義模型的方案，突破了這條路線的技術(shù)瓶頸。在未來的研究中，我們計劃從以下幾個方面進(jìn)一步推進(jìn)：

• 對模型剪枝和知識壓縮方案的進(jìn)一步探索；
• 利用技術(shù)路線中可自定義模型結(jié)構(gòu)的優(yōu)勢，探索更高效的訓(xùn)練、推理算法方案；
• 更廣泛的開源和社區(qū)合作，我們將繼續(xù)推進(jìn)模型和方案的開放性，并與學(xué)術(shù)界和工業(yè)界合作，共同構(gòu)建更強大的代碼生成系統(tǒng)。

引用

[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.