CodeGeeX: 多语言代码生成模型

CodeGeeX: 多语言代码生成模型 2022年9月19日 | 语言：En

CodeGeeX是一个具有130亿参数的多编程语言代码生成预训练模型。CodeGeeX采用华为MindSpore框架实现，在鹏城实验室“鹏城云脑II”中的192个节点（共1536个国产昇腾910 AI处理器）上训练而成。截至2022年6月22日，CodeGeeX历时两个月在20多种编程语言的代码语料库（>8500亿Token）上预训练得到。CodeGeeX有以下特点：

全新多编程语言评测基准HumanEval-X：HumanEval-X是第一个支持功能正确性评测的多语言、多任务的基准，包含820个人工编写的高质量代码生成题目、测试用例与参考答案，覆盖5种编程语言（Python、C++、Java、JavaScript、Go），支持代码生成与代码翻译能力的评测。

在HumanEval-X代码生成任务上，与其它开源基线模型相比，CodeGeeX取得了最佳的平均性能。

近年来，使用代码语料训练的大规模预训练模型取得飞速进步。Codex通过使用Python求解初级编程问题，展示了预训练模型在该方面的潜力。此后，一系列代码生成模型亦得以面世，比如AlphaCode、CodeGen、InCoder、PolyCoder、PaLMCoder等。这些模型都使用了多种编程语言进行训练，但是它们往往仅在Python上做正确性评测，在其它语言上的生成性能尚不明确。

现有的公开评测基准主要关注两种评价指标：字符串相似性（string similarity）或功能正确性（functional correctness）。第一种指标，如CodeXGLUE和XLCoST多语言基准，涵盖了代码补全、翻译、概括等任务。它们使用了BLEU和CodeBLEU这类判断相似性的指标，但这些指标并不能很好反映代码是否正确。相反地，第二种指标通过测试用例来判断代码功能上是否正确，如HumanEval、MBPP、APPS等基准。然而，这些基准只支持Python，并不支持其他编程语言。缺乏评价代码正确性的多语言基准，阻碍了多语言代码生成模型的发展。

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

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

语料：CodeGeeX的训练语料由两部分组成。第一部分是开源代码数据集，The Pile 与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相同的分词器，并将空格处理为特殊标识符，词表大小为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优化器降低显存占用。

在开发与训练过程中，我们和华为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整体训练效率有极大提升