KwaiCoder-23BA4-v1 是快手 Kwaipilot 团队最新开源的自研代码续写大模型,模型的训练依托于 Kwaipilot 团队提出的一种高效的模型训练方案。通过结合模型剪枝、知识蒸馏、细粒度合并等技术,相比传统方法以 1/30 的成本完成了 23B 宽 MoE 架构代码续写模型的训练,并在多个代码领域评测集上实现了新的 SOTA。
开源地址: huggingface.co/Kwaipilot/K...
我们的路径:
- 高质量数据:我们强调筛选和整理高质量代码、文本与合成数据,而非单纯扩大数据集规模,从而确保数据的相关性和一致性;
- 一种经济且高效的模型训练路线:通过结合模型知识蒸馏与细粒度合并,完成模型的知识压缩(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.