亚马逊AWS官方博客

云学堂:业务代码能力提升 44%,基于 Amazon SageMaker 大模型微调赋能代码生成的创新实践

关于云学堂

云学堂(Nasdaq: YXT)作为一家科技企业,为客户提供数字化企业学习解决方案服务,其中包括:软件 SaaS 服务、内容解决方案服务、数字化企业学习运营服务,是企业培训领域的创新型公司。2024 年云学堂成功在美上市。

自 2011 年成立以来,云学堂已服务包括世界 500 强、中国 500 强在内的众多行业头部公司,获得业界一致认可与好评。

背景描述

随着软件开发复杂性的增加和对效率的不断追求,云学堂意识到传统的代码编写方式已经无法满足快速迭代和高质量交付的需求。借鉴行业内的成功经验,结合生成式 AI 技术的发展,云学堂计划通过现有的大量的 Java 业务代码数据,结合先进的大语言模型技术技术,开发一个能够智能生成和补全代码的大模型。该模型将专注于理解和生成 Java 后端代码,帮助开发者快速编写高质量代码。它不仅能够自动生成代码片段,还能根据上下文智能补全代码,并提供代码优化建议。通过这种方式,提高开发效率,减少人为错误,并推动创新。

关于 SageMaker

Amazon SageMaker 是一个全方位的机器学习服务,专为快速构建、训练和部署机器学习模型而设计,特别适用于大规模模型的训练。SageMaker 提供集成的数据准备工具,支持从 S3、Redshift 等数据源导入数据,并内置多种机器学习算法和深度学习框架,如 TensorFlow 和 PyTorch。其强大的分布式训练功能允许在大规模数据集和复杂模型上进行训练。此外,SageMaker 提供简便的模型部署功能,支持一键部署到托管的实时端点。通过这些功能,SageMaker 显著简化了机器学习工作流,使开发人员和数据科学家能够专注于模型创新,而无需担心底层基础设施。

训练方案

对于整体的训练方案,构建了基于 SageMaker 的训练流水线。 首先,通过 SageMaker Notebook 对业务代码库进行 AST 解析,生成 FIM 数据集,提升代码补全能力。接着,使用 SageMaker Training Job 对基础模型进行微调。模型评估阶段包括泛化和 FIM 测试,确保模型性能。

数据准备阶段

为了训练一个强大的代码生成模型,我们的团队投入了大量时间和精力来准备高质量的训练数据。我们基于业务代码库,使用 AST(抽象语法树)技术提取了业务代码,并通过该技术生成了 FIM(Fill-in-the-Middle)数据,这种数据结构有助于模型在代码补全和插入任务中的表现。同时,我们对原始代码进行了随机分割,以进一步提高模型在代码补全方面的泛化能力。最终,我们构建了一个包含海量 Java 代码语料库的数据集,总量超过亿级 tokens。其中,67% 为业务代码数据,20% 来自 GitHub 的 Java 代码,13% 为中英文自然语言数据,以确保模型能够在理解代码和自然语言对应关系方面表现优异。

训练数据格式

一般来讲,训练一个代码模型会主要使用两种训练任务来提升模型在下游应用的表现,分别是代码补全(Code Completion)以及代码插入(Code Inffing)。

  • 代码补全

代码补全是最常见的一种代码模型训练任务,模型需要根据给定的代码上下文,预测下一个 token(单词或子词)。

以使用企业内部的代码数据继续预训练的方式为例,训练数据通常是一个完成的样本,让模型将代码的生成能力迁移到企业内部的代码逻辑或者企业内部库的 API 上。如果是 SFT 阶段,训练数据通常是将一个代码片段切分成前缀(上下文)和目标对,即给定一个训练样本的前半部份, label 则是该样本的原始后半部分。

  • 代码插入

代码插入任务的目标是让模型学习预测被掩码(masked)或被噪声污染的代码片段中的缺失部分。具体来说,训练过程包括以下步骤:

  1. 在完整的代码中选择一些 token 位置,用特殊的掩码 token(如[MASK])替换
  2. 将带有掩码的代码片段输入到模型
  3. 模型的目标是预测那些被掩码位置原本的 token

代码插入任务有以下几个好处:

  • 无需额外的监督数据,可以直接利用原始代码库
  • 通过随机采样/AST 分析后进行采样和掩码,可以提升训练数据多样性
  • 模型需要学习捕捉代码的语法和语义信息才能正确预测,与纯序列生成相比,需要具有更强的语言建模能力

在本次的代码插入训练任务,采用了 PSM (Prefix – Suffix – Middle)这个形式组织训练数据,该任务数据按照以下格式处理,更多详细信息可以参考 DeepSeek Coder 的论文。

全参数微调阶段

在完成数据准备之后,我们进入了模型训练的核心阶段——全参数微调。我们选择了当前热门的语言模型 DeepSeek Coder 6.7B 作为基座模型,并在此基础上使用优化过的数据语料进行训练。为了提升模型的性能和适应性,我们在微调过程中综合运用了多种技术,包括 PiSSA(Position-aware Siamese Self-Attention)技术、LoRA(Low-Rank Adaptation)以及全参数微调(Full Fine-Tuning)。

在训练过程中,我们密切关注了模型的训练收敛情况。从训练曲线中可以看到,模型的训练损失(training loss)和评估损失(eval loss)都呈现出逐步下降的趋势,并且随着训练步数的增加,两者都逐渐趋于平稳。这表明模型在经过充分的训练后,逐渐达到了收敛状态。最初的训练损失相对较高,但随着训练的进行,损失迅速下降,并在大约 2000 步左右后开始趋于稳定。随着进一步的训练,模型的损失值进一步下降,最后达到了一个较低且平稳的水平。

此外,为了有效处理大规模模型的训练,并提高训练效率和缩短训练时间,我们采用了 DeepSpeed 的分布式训练框架。 在 Amazon SageMaker 平台上,我们通过使用 Training Job 进行大规模分布式训练,充分发挥了这些技术的优势。最终,我们成功地开发出了一个功能更加完善、性能更加强大的 DeepSeek-Coder-6.7b-Base 代码生成模型。这个新版本的模型在生成能力上有了显著提升,并且训练过程中的良好收敛性为模型的稳定性和高效性提供了坚实的基础。

训练参数

DeepSpeed zero_type:3 ,为了达到最佳的训练效能。

learning rate:2.0e-5,设定的学习率为 0.00002。这个参数决定了模型在每一步更新时权重调整的步幅大小。

warmup ratio:0.1,这意味着训练过程中,前 10% 的步数将用于从零逐渐增加学习率,确保模型在训练初期的稳定性。

微调中的 train/eval loss

模型评估阶段

在完成 DeepSeek Coder 模型的微调后,我们采用了全面的评估方法,对模型进行了深入的分析,以确保评估结果的客观性和准确性。

首先,我们在业务 FIM 能力上进行了重点评估。结果显示,微调后的模型在多个关键指标上表现显著提升。具体来说,CodeBleu 得分从基础模型的 38.96 提升至 56.07,ngram 匹配分数从 15.32 上升至 31.86,weighted ngram 匹配分数从 17.60 增至 40.62,而 syntax 匹配得分也从 47.14 跃升至 65.06。这些数据表明,微调后模型在处理业务代码生成任务时,表现出更强的理解和生成能力,明显优于基础模型。

同时,我们也对模型的基础能力进行了评估。虽然在基础任务(如 Recode、HumanEval 和 MBPP)中,微调后的模型表现略有下降,Recode 得分从 42.13 下降至 39.06,HumanEval 得分从 44 下降至 36.8,MBPP 得分从 47.6 下降至 41.86,但这种下降幅度在可接受的范围内,并且这种基础能力的轻微下降,是以换取业务 FIM 能力的大幅提升为代价的。

通过此次评估,可以看出微调后的模型在业务 FIM 能力上有了显著的进步,尽管基础能力有所下降,但这种变化是可以接受的,且符合预期。这种能力上的此消彼长现象,正是为了在特定业务场景中提升模型性能所做的策略性选择。因此,我们可以认为,微调后的模型在不显著牺牲基础代码能力的情况下,实现了业务相关任务上的显著提升,成功达到了项目的预期目标。

业务 FIM 能力测评(提升 44%)

模型/评测维度 CodeBleu ngram_match_score weighted_ngram_match_score syntax_match_score dataflow_match_score
deepseek-coder-6.7b-base 38.96 15.32 17.60 47.14 29.76
deepseek-coder-6.7b-Fine-tuning 56.07 31.86 40.62 65.06 33.75

基础能力测评

模型/评测维度 Recode HumanEval MBPP
deepseek-coder-6.7b-Base 42.13 44 47.6
deepseek-coder-6.7b-Fine-tuning 39.06 36.8 41.86

测效果(使用 IntelliJ IDEA 2024.2.1)

总结

本文详细介绍了云学堂利用 Amazon SageMaker 平台,针对 Java 代码生成领域进行大语言模型微调的探索和实践。通过精心的数据准备、创新的训练技术和全面的模型评估,云学堂成功开发出基于 DeepSeek-Coder 性能卓越的业务代码生成模型。该模型在处理业务代码生成任务时表现出色,在保证基础代码能力的同时,实现了业务场景下的显著提升。未来,云学堂将继续携手亚马逊云科技,探索前沿 AI 技术在企业数字化领域的应用。

本篇作者

朱可

云学堂高级工程师,专注于前沿人工智能技术的探索与工程化应用,熟悉 Transformer 核心架构,对模型训练有丰富经验,致力于开发安全、高效、且易于维护的 AI 工程化解决方案。

曹镏

亚马逊云科技解决方案架构师,专注于为企业级客户提供信息化以及生成式 AI 方案的咨询与设计,在 AI/ML 领域具有解决实际问题能力以及落地大模型训练项目的经验。

董孝群

亚马逊云科技解决方案架构师,GCR GenAI SSA,负责生成式 AI 解决方案的设计,曾在百度,粤港澳大湾区数字经济研究院供职,在 nlp 领域有着丰富经验,目前专注于大语言模型相关解决方案的研发。