CodeGeeXr多语言代码生成模型

[toc]

背景

去年copilot微软发布的基于AI的辅助工具因为利用Github开源代码，造成了不小的冲击，不过copilot的功能还是让人印象深刻

代码补全：很多ide都支持，单copilot 增强了 IDE 的补全功能。

注释即代码：通过编写注释然后一路根据 copilot 的提示编写出完整的程序

copilot虽好，但是随着收费，热度也在下降，而最近清华大学和华为推出的开源工具CodeGeeX，不仅仅兼容copilot功能，还支持了代码翻译等功能，根据宣传，主要有一下几个特点：

高精度代码生成：支持生成Python、C++、Java、JavaScript和Go等多种主流编程语言的代码，在HumanEval-X代码生成任务上取得47%~60%求解率，较其他开源基线模型有更佳的平均性能。

跨语言代码翻译：支持代码片段在不同编程语言间进行自动翻译转换，翻译结果正确率高，在HumanEval-X代码翻译任务上超越了其它基线模型。

自动编程插件：CodeGeeX插件现已上架VSCode插件市场（完全免费），用户可以通过其强大的少样本生成能力，自定义代码生成风格和能力，更好辅助代码编写

相关地址

代码仓库：https://github.com/THUDM/CodeGeeX

英文官网：http://keg.cs.tsinghua.edu.cn/codegeex/

中文官网：https://keg.cs.tsinghua.edu.cn/codegeex/index_zh.html

代码生成demo：https://models.aminer.cn/codegeex/zh-CN/playground

代码翻译：https://models.aminer.cn/codegeex/zh-CN/codeTranslator

vscode插件下载地址：https://marketplace.visualstudio.com/items?itemName=aminer.codegeex

CodeGeeX: 多语言代码生成模型

下面的内容从官网转载：

架构：CodeGeeX是一个基于transformers的大规模预训练编程语言模型。它是一个从左到右生成的自回归解码器，将代码或自然语言标识符（token）作为输入，预测下一个标识符的概率分布。CodeGeeX含有40个transformer层，每层自注意力块的隐藏层维数为5120，前馈层维数为20480，总参数量为130亿。模型支持的最大序列长度为2048。

左侧：CodeGeeX训练数据中各编程语言占比。 右侧：CodeGeeX训练损失函数随训练步数下降曲线。

语料：CodeGeeX的训练语料由两部分组成。第一部分是开源代码数据集，The Pile [13]与CodeParrot。The Pile包含GitHub上拥有超过100颗星的一部分开源仓库，我们从中选取了23种编程语言的代码。第二部分是补充数据，直接从GitHub开源仓库中爬取Python、Java、C++代码；为了获取高质量数据，我们根据以下准则选取代码仓库：1)至少拥有1颗星；2)总大小<10MB；3)不在此前的开源代码数据集中。我们还去掉了符合下列任一条件的文件：1)平均每行长度大于100字符；2)由自动生成得到；3)含有的字母不足字母表内的40%；4)大于100KB或小于1KB。为了让模型区分不同语言，我们在每个样本的开头加上一个前缀，其形式为[注释符] language: [语言]，例如：# language: Python。我们使用与GPT-2[14]相同的分词器，并将空格处理为特殊标识符，词表大小为50400。整个代码语料含有23种编程语言、总计1587亿个标识符（不含填充符）。

国产平台实现与训练

我们在Mindspore 1.7框架上实现了CodeGeeX模型，并使用鹏城实验室的全国产计算平台上进行训练。具体来说，CodeGeeX使用了其一个计算集群中的1536个昇腾910 AI处理器（32GB）进行了两个月左右的训练（2022年4月18日至6月22日）。除了Layer-norm与Softmax使用FP32格式以获得更高的精度与稳定性，模型参数整体使用FP16格式，最终整个模型需要占用约27GB显存。为了增加训练效率，我们使用8路模型并行和192路数据并行的训练策略，微批大小为16、全局批大小为3072，并采用ZeRO-2优化器[15]降低显存占用。

在开发与训练过程中，我们和华为Mindspore团队合作，对MindSpore框架进行了部分优化，进而大幅度提升训练效率。比如，我们发现矩阵乘法的计算时间占比仅为22.9%，大量时间被用于各类其它算子，因此实现了一系列算子融合，包括单元素算子融合、层归一化算子融合、FastGelu与矩阵乘法融合、批量矩阵乘法与加法融合等；再比如我们还对矩阵乘法算子的维度实现自动搜索调优，使其搜索出效率最高的计算维度组合。这些优化为训练速度带来了显著提升，在同等GPU卡数规模下（128卡），昇腾910对CodeGeeX这一模型的训练效率从约为NVIDIA A100的16.7%提升至43%；在千卡规模下，昇腾910训练效率相比自身优化前提升近300%。下表为优化前与优化后，昇腾910与英伟达A100在CodeGeeX模型训练上的效率对比。使用优化后的软硬件训练时，CodeGeeX单日训练量可达到543亿个标识符（含填充符）。我们还尝试进一步加入流水线优化（训练时未实装），整体训练效率再次提升，单日训练量达到841亿个标识符（含填充符），证明了国产深度学习平台与工具的快速迭代能力以及强大竞争力。

Ascend 910 与 NVIDIA A100优化前后训练效率对比

进一步加入流水线并行等优化后，Ascend 910整体训练效率有极大提升

除了对Ascend平台的支持，我们也正在将模型迁移适配到其它平台，并将在近期开源跨平台代码和模型权重，让研究者可以在不同平台上运行CodeGeeX。

HumanEval-X: 多语言代码生成基准

为了更好地评测代码生成模型的多语言生成能力，我们构建了一个新基准HumanEval-X。此前，多语言代码生成能力是基于语义相似度（比如CodeBLEU[10:1]）衡量的，具有一定误导性；HumanEval-X则可用于衡量生成代码的功能正确性。HumanEval-X包含820个高质量手写样本，覆盖Python、C++、Java、JavaScript、Go，可用于多种任务。

HumanEval-X支持的任务示例。声明、描述、解答分别用红、绿、蓝色标注。代码生成将声明与描述作为输入，输出解答。代码翻译将两种语言的声明与源语言的解答作为输入，输出目标语言的解答。

HumanEval-X中每个语言的样本，包含了声明、描述和解答，它们之间的组合可以支持不同的下游任务，包括生成、翻译、概括等。我们目前关注两个任务：代码生成与代码翻译。对于代码生成任务，模型将函数声明与文档字符串作为输入，输出函数实现；对于代码翻译任务，模型将两种语言的函数声明与源语言的实现作为输入，输出目标语言上的实现。我们在代码翻译任务中不将文档字符串输入模型，以避免模型直接通过描述生成答案。在两种任务下，我们都采用Codex[1:2]所使用的无偏pass@k指标，判断生成代码的功能正确性。

多语言代码生成

左上：在HumanEval-X的代码生成任务中，模型在所有语言上的平均表现。其他：在五种语言上具体的pass@k（k=1,10,100）性能。CodeGeeX的平均表现优于InCoder-6.7B和CodeGen-Multi-6B/16B。

我们将CodeGeeX与另外两个开源代码生成模型进行比较，分别为Meta的InCoder与Salesforce的CodeGen，选取InCoder-6.7B、CodeGen-Multi-6B 与 CodeGen-Multi-16B。CodeGeeX能获得最佳的平均性能，显著超越了参数量更小的模型(7.5%~16.3%的提升)，与参数量更大的模型CodeGen-Multi-16B表现相当（平均性能 54.76% vs. 54.39%）。除此之外，我们还探索了将生成次数分配给不同语言的效果，在按照训练数据比例分配生成次数时，CodeGeeX的正确率比其在任一单语言下的正确率都更高（如红色虚线圈所示）。

多语言生成超越最佳单一语言生成

在HumanEval-X中，CodeGeeX在5种语言下对各问题的解答率，按Python下的解答率排序。

我们观察了每道题目的具体解答情况，发现题目在不同语言上的解答率有着较大区别，某些题目更适合用特定语言进行解答。这个现象启发我们将固定的生成次数分配到多种语言上，从而让生成结果更多样，增大其中至少有一个正确解答的概率。

固定生成预算时HumanEval-X下的多语言生成结果。每种方法最佳模型的结果用粗体显示，每个模型最佳方法的结果用斜体显示。

如果采用平均分配，比如使用每种语言生成20次，即可让正确率有所提升。我们还使用另一种启发式方法，根据训练语料中各语言占比来分配生成的次数，可以将CodeGeeX的pass@100 从60.92%提升到63.92%。使用这种启发式方法后，我们测试的所有模型的性能都有所提升，pass@10与pass@100相对最好的单语言生成提升了1~4%。在上述的分配方式中，CodeGeeX的表现都比其它基线模型更好。

使用上述加权分配策略时，CodeGeeX显示出优秀的多语言能力，优于参数量更大的CodeGen-Multi-16B。同时解答率随着采样数量的增加可以进一步提高。

跨语言代码翻译

HumanEval-X上的代码翻译任务结果。加粗结果表示在每种语言pass@k上的最佳效果。

我们还评测了模型在多语言间代码翻译上的性能。对于CodeGeeX，我们评测了未经微调的CodeGeeX-13B与经过微调的CodeGeeX-13B-FT（使用XLCoST [8:1]中代码翻译任务的训练集与一部分Go语言数据微调）。如上表显示，模型对特定语言存在偏好，比如CodeGeeX擅长将其他语言翻译为Python与C++，而CodeGen-Multi-16B擅长翻译为JavaScript和Go，这可能是由于训练集中的语料占比存在差异。在20个翻译对中，我们还观察到两种语言互相翻译的表现常常是呈负相关的，这可能说明现有的模型还不足以学好所有的语言。

在线生成与翻译DEMO

我们为上述两个任务开发了DEMO：代码生成和代码翻译，欢迎点击体验！

自动编程VS Code插件

基于CodeGeeX，我们开发了一个免费的VS Code插件（插件市场搜索“CodeGeeX”即可下载），辅助多语言编程开发。除了多语言代码生成/翻译能力，我们还利用CodeGeeX强大的少样本生成能力，将其变成一个可自定义的编程助手。少样本生成意味着CodeGeeX不需要在特定任务上进行微调，用户只需通过提供自定义提示语句，即可实现一些有趣的功能。比如，用户可以在提示中加入具有特定代码风格的代码，CodeGeeX会加以模仿，生成类似风格的代码。用户也可以加入一些特定格式的代码让CodeGeeX掌握新技能，如逐行解释代码（如动图所示）。马上前往VS Code尝试CodeGeeX，自定义你的编程助手吧！

上述例子中，我们在输入中提供了额外提示—-一段逐行解释代码的示例，CodeGeeX加以模仿，将现有的一段代码按照相同的方式进行解释。可以通过类似的方式，加入其他的提示，让模型掌握更多的能力.

未来研究

在之前工作的基础上，CodeGeeX进一步展示了大规模预训练模型在代码生成上的强大能力，以及进一步扩展到各种形式化语言上的潜力。这仅仅是第一步，还有很多问题值得我们在未来继续探索，这里给出三点我们的思考。

第一，我们发现模型容量对多语言能力十分关键。人类程序员在掌握几种编程语言的情况下，能快速迁移到其他语言。与之相反，模型则需要为每种语言分配较大的模型容量来存储相关的知识（一个典型的例子是有着5400亿参数量的PaLMCoder，依靠超大的模型容量取得了惊人的效果），却没有很好抽象出编程语言的高阶知识。如何让模型学习到编程语言的共性，从而快速迁移到其他语言上仍是一个巨大的挑战。

第二，CodeGeeX展现了一定的推理能力，但并没有推广到所有语言上。CodeGeeX能使用不同语言解决编程问题，但各种语言擅长解答的问题分布却有较大差异。尽管模型已经能学习到各种语言正确的语法，却无法用不同语言解决同一个问题。我们猜想这可能与不同编程语言的特性有关（比如Python库功能更加完善），或者是因为在训练集中出现了某种语言的类似实现，但缺少了其他语言的。无论是哪种情况，都反映了当前的模型在推理能力上存在局限，让模型获得更加可靠的推理能力还需要长期的研究。

第三，CodeGeeX的少样本生成能力亟待进一步探索。对于大规模预训练模型，微调（fine-tuning）的成本是昂贵的，如何使用少量的样本就让模型生成想要的代码，对代码生成模型的实用化具有非常大的意义。近期的一些工作，如思考过程提示方法（chain-of-thought prompting）[16]就展现了惊人的结果，这极大启发了我们将这一技术也用于CodeGeeX中。这同时也是我们设计自定义编程助手的初衷，希望程序员和研究者们能和我们一起探索如何更好地使用CodeGeeX。