paper

CodeGeeX

1.abstract

  • 130亿个参数的多语言代码生成模型CodeGeeX在23种编程语言的8500亿个token上进行了预训练,代码生成和翻译任务上都优于类似规模的多语言代码模型。
  • 在ide上有扩展插件
  • 开源地址：https://github.com/THUDM/CodeGeeX

2.introduction

通过将程序视为语言序列，神经序列架构，如rnn和transformer(Vaswani et al.， 2017)，可以自然地应用于代码生成。事实上，基于transformer的技术(Svyatkovskiy等人，2020;Sun et al.， 2020)已经展示了自动程序编写的潜力，开始生成既语法正确又有效的代码,当大型语言模型(具有数十亿参数的转换器)遇到大量开源代码数据时，这一进展得到了显著的推进。

内容：

  • 开源

  • CodeGeeX还支持在Ascend和NVIDIA gpu上的跨平台推理。

  • CodeGeeX除了像Codex那样生成代码和完成代码之外，还支持语言对之间的代码解释和代码翻译任务

  • 与同类规模的知名多语言代码生成模型(包括CodeGen-16B、gpt - nex - 20b、InCode-6.7B和GPT-J-6B)相比，它具有一致的性能优势(参见图1 (b)和©)。

  • IDE中构建了免费的CodeGeeX扩展，目前包括Visual Studio Code、JetBrains和腾讯云工作室(一个Web IDE)。

  • 评估多语言代码模型的HumanEval- x基准，如下:1)HumanEval (Chen等人，2021)-由OpenAI开发用于评估codex -和其他基准(Austin等人，2021;Hendrycks等人，2021;Nijkamp等人，2022)只包含单一语言的编程问题和2)现有的多语言数据集(Ren等人，2020;Lu et al.， 2021;Zhu et al.， 2022)使用BLEU (Papineni et al.， 2002)等字符串相似度指标进行评估，而不是真正验证生成代码的功能正确性。

本工作的贡献可以概括如下:

  • 我们开发并发布CodeGeeX，这是一个13B预训练的23种语言代码生成模型，在相同规模的多语言基线上，它在代码生成和翻译方面表现出一致的优异表现。
  • We build the CodeGeeX extensionson VSCode4、JebBrains5 andTencentCloudStudio。与Copilot相比，它支持更多样化的功能，包括代码补全、生成、翻译和解释。根据用户调查，CodeGeeX可以提高83.4%的用户的编码效率。
  • 我们手工制作HumanEval-X基准，以评估代码生成和翻译任务的多语言代码模型的功能正确性，促进对预训练(多语言)代码模型的理解和开发。

3.CodeGeeX Model

CodeGeeX是一个多语言代码生成模型，具有130亿(13B)个参数，在23种编程语言的大型代码语料库上进行预训练。截至2022年6月22日，CodeGeeX已经在1,536个Ascend 910 AI处理器集群上对超过8500亿个代币进行了两个多月的训练。

3.1 CodeGeeX结构：

transformer的主干CodeGeeX遵循生成式预训练(GPT)架构(Radford等人，2018)，采用自回归(编程)语言建模的解码器风格。CodeGeeX的核心架构是一个39层的transformer解码器。在每个transformer层(见图2)中，我们应用了多头自注意力机制(Vaswani等人，2017)，然后是MLP层，以及层规范化(Ba等人，2016)和剩余连接(He等人，2016)。我们使用了一种近似的GELU(高斯线性单元)操作(Hendrycks和Gimpel, 2016)，即FastGELU，它在Ascend 910 AI处理器下更高效。

生成式预训练目标。通过采用GPT范式(Radford et al.， 2019;Chen et al.， 2021)，我们在大量未标记的代码数据上训练模型。其原理是迭代地将代码标记作为输入，预测下一个标记，并将其与基本事实进行比较。具体来说，对于任意输入序列{x1, x2，…，长度为n的xn}， CodeGeeX的输出是下一个token P(xn+1|x1,x2，…，xn，Θ) = pn+1∈[0,1]1×v的概率分布，其中Θ表示模型的所有参数，v表示词汇表大小。通过将其与实分布进行比较，即ground-truth令牌的一个单热向量yn+1∈{0,1}1×v，我们可以优化累积交叉熵损失

顶层查询层和解码。原始的GPT模型使用池函数来获得最终输出。我们在所有其他变压器层之上使用额外的查询层(Zeng et al.， 2021)，通过关注获得最终嵌入。如图2所示，顶层查询层的输入用位置n + 1的查询嵌入代替查询输入Xin。最后的输出乘以词嵌入矩阵的转置得到输出概率。对于解码策略，CodeGeeX支持贪婪、温度采样、top-k采样、top-p采样和波束搜索。最后，去标记化将把选定的令牌ID转换为实际的单词。

3.2 预训练的设置

代码语料库。训练语料库包含两个部分。第一部分来自开源代码数据集。第二部分是直接从GitHub公共存储库中抓取的Python、Java和c++的补充数据，这些数据在第一部分中没有出现。

标记。第一步是将代码片段转换为数值向量。考虑到1)代码数据中存在大量自然语言注释，2)变量、函数和类的命名通常是有意义的单词，我们将代码数据视为文本数据并应用GPT-2标记器(Radford et al.， 2019)。它是一个BPE(字节对编码)(Sennrich等人，2015)标记器，它使用固定大小的词汇表和可变长度的字符来处理开放词汇表问题。初始词汇表大小为50,000，我们按照Chen等人(2021)的方法将多个空格编码为额外的标记，以提高编码效率。具体来说，L空格由<|extratoken_X|>表示，其中X=8+L。由于词汇表包含来自各种自然语言的标记，因此它允许CodeGeeX处理除英语以外的其他语言的标记，如中文、法语、俄语、日语等。最终的词汇量是v = 52,224。在标记化之后，任何代码片段或文本描述都可以转换为整数向量。详情见附录A.2。

输入词和位置嵌入。给定了标记，下一步是将每个标记与单词嵌入关联起来。通过在单词嵌入矩阵Wword∈Rv×h中查找令牌ID，其中h = 5120为隐藏大小，为每个令牌获得一个可学习的嵌入xword∈Rh。为了获取位置信息，我们还采用了可学习的位置嵌入，将当前位置ID映射到一个可学习的嵌入xpos∈Rh，从Wpos∈Rnmax×h，其中nmax = 2048为最大序列长度。然后，将两个嵌入相加，得到模型的输入嵌入xin = xword + xpos。最后将整个序列转化为输入嵌入Xin∈Rn×h，其中n为输入序列长度。

3.3 CodeGeeX训练

平行训练提升910。CodeGeeX在使用Mindspore (v1.7.0)的Ascend 910 AI处理器(32GB)集群上进行训练。在预训练期间，我们面临并解决了许多未知的技术和工程挑战，因为与NVIDIA gpu和PyTorch/TensorFlow相比，Ascend和Mindspore相对较新。整个预训练过程在192个节点和1536个AI处理器上花费了两个月的时间，在此期间，模型消耗了850B个令牌，相当于5+个epoch(213,000步)。详细的配置见表2。

为了提高训练效率，我们采用8路模型并行训练和192路数据并行训练，启用ZeRO-2 (Rajbhandari et al.， 2020)优化器，进一步减少优化器状态的内存消耗。最后，微批大小为每个节点16个，全局批大小达到3,072个。

具体来说，我们使用Adam优化器(Kingma and Ba, 2014)来优化方程2中的损失。模型权值采用FP16格式，但为了更高的精度和稳定性，我们使用FP32进行层规范和softmax。该型号的GPU内存约为27GB。我们从初始学习率1e-4开始，并应用余弦学习率衰减。

在为期两个月的培训中，随着培训的进行，CodeGeeX的培训损失不断减少。我们评估了HumanEval-X代码生成任务上的检查点，并观察到性能在不断提高。详见附录A.3中的图13和图14。

培训效率优化。在培训过程中，我们积极尝试优化Mindspore框架，以释放Ascend 910的力量。值得注意的是，我们采用了以下技术，显著提高了培训效率:

  • 内核融合:我们融合了几个元素明智的运算符来提高Ascend 910的计算效率，包括Bias+LayerNorm, BatchMatmul+Add, FastGeLU+Matmul, Softmax等。我们还优化了LayerNorm运算符以支持多核计算。
  • 自动调整优化:加载模型时，Mindspore首先将它们编译为静态计算图。它使用Auto Tune工具来优化运算符的选择(例如，不同维度的矩阵乘法)。并应用图优化技术处理算子融合和常数折叠问题。

表3是我们优化前后的训练效率对比。总体效率是通过每天训练的代币来衡量的。我们观察到，与未优化的实现相比，每个处理器的效率提高了3倍，1,536个gpu的总体令牌吞吐量提高了224%。

3.4 快速推理

为了服务于预训练的CodeGeeX，我们实现了一个纯PyTorch版本的CodeGeeX，它支持在NVIDIA gpu上进行推理。为了实现快速和高效的推理，我们将量化和加速技术应用于预训练的CodeGeeX。

量化。我们应用训练后量化技术来减少CodeGeeX在推理过程中的内存消耗。在FP16到INT8的所有线性变换中，我们使用常见的绝对最大量化对权重W进行变换

加速。量化后，我们使用NVIDIA FasterTransformer (FastTrans)进一步实现更快版本的CodeGeeX。它通过使用层融合、GEMM自动调优和硬件加速功能支持高度优化的操作。对于INT8量化版本，我们还实现了一个自定义内核，用于加速INT8权值与FP16激活向量之间的混合精度矩阵乘法。根据表4,INT8量化加上FastTrans实现在单个GPU上实现了最快的推理速度和最低的GPU内存消耗。每个令牌的推理时间在13ms以内(1.61秒/ 128个令牌)。我们还比较了LLM.int() (Dettmers等人，2022)和Oneflow (Yuan等人，2021)中的推理速度。

4.The HumanEval-X Benchmark

我们开发了用于评估多语言代码模型的HumanEval-X基准。针对五种主要语言(c++、Java、JavaScript、Go和Python)定义了164个代码问题，导致164×5=820个问题-解决方案对。对于每个问题，它都支持代码生成和代码转换。这些问题的示例可以在附录A.5中找到。

5.结论

我们引入了CodeGeeX，一个13B预训练的23语言代码生成模型，并构建了HumanEval-X，以填补多语言代码生成的空白。在代码生成和翻译任务上，CodeGeeX始终优于相同规模的开源多语言基线。建立在CodeGeeX上的扩展在提高编码效率方面带来了显著的好处。CodeGeeX的多语言性带来了用一组无处不在的形式化语言解决问题的潜力。我们开源CodeGeeX的目的是帮助研究人员和开发人员广泛地利用大型预训练模型来生成代码。
CodeGeeX的多语言能力显示了用一组无处不在的形式化语言解决问题的潜力。在这里，我们分享三个我们的观察作为未来的方向。

首先，我们发现模型能力对多语言编程能力至关重要。该模型从学习多种语言中获益并不是微不足道的。人类程序员可以抽象编程的高级概念，因此学习一种语言可以帮助他们掌握其他语言。相反，该模型似乎需要很大的容量来同时存储每种语言的知识。如何帮助模型提取最基本的编程知识仍然是一个研究挑战。

其次，与其他类似，CodeGeeX显示了作为模型的推理潜力，尽管它缺乏很强的通用性。我们证明了CodeGeeX可以解决不同语言的问题。然而，不同语言之间的通过率分布差异很大，即有时使用不同的语言无法解决相同的问题。我们假设这可能与某些特定于语言的特性有关(例如，有些问题在Python中更容易解决)，或者它可能仅仅是由于在训练数据中出现了类似的特定于语言的实现。无论哪种情况，模型要具有可靠的推理能力还有很长的路要走。

最后，CodeGeeX的少拍能力值得探索。而不是使用昂贵的微调方法，我们可以使用几个例子来启动，使模型达到可比较的性能。最近的工作，如思维链(CoT)提示(Wei et al.， 2022)使用这种方法显示了令人印象深刻的结果，激励我们在代码模型中检查CoT。