《CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Evaluations on HumanEval-X》

论文地址：https://arxiv.org/pdf/2303.17568.pdf

代码地址：https://github.com/THUDM/CodeGe

一、简介

CodeGeeX，是一个具有130亿个参数的多语言模型，用于代码生成，在23种编程语言的8500亿个token上进行了预训练。

CodeGeeX的特点

(1)CodeGeeX模型本身以及训练代码都是开源的，有助于理解和改进预训练代码模型。也支持在Ascend和NVIDIAGPUs等不同平台上推理。

(2)除了代码生成和代码补全，也支持代码解释和代码翻译。

(3)与知名的代码生成模型相比(CodeGen-16B、GPT-NeoX-20B、InCode-6.7B和GPT-J-6B)，CodeGeeX的表现优于其他模型。

二、模型

CodeGeeX是一个具有13B个参数的代码生成模型，由39层从左到右的transformer decoders和一个top query layer组成。它将文本/代码tokens作为输入，并自回归地输出下一个标记的概率

1、使用Transformer为主干。

与GPT-3Codex类似，CodeGeeX也是生成式预训练(GPT)架构，使用纯解码器的GPT架构，采用自回归(编程)语言建模。

CodeGeeX的核心架构是一个39层的transformer解码器。在每个transformer层中，应用了多头自注意力机制，然后是MLP层，以及层归一化和残差连接。使用类GELU的FastGELU，它在Ascend910AI处理器下效率更高:

2、生成式预训练。

通过采用GPT范式，在大量未标记的代码数据上训练模型。其原理是迭代地将代码token作为输入，预测下一个token，预测下一个token并与真实的token进行比较。对于长度为n的任意输入序列 {x1,...xn} ，CodeGeeX的输出都是下一个token的概率分布

Θ 表示所有参数，v是词表大小。通过将预测token与真实分布进行比较，可以优化交叉熵损失函数：

3、Top Query层和解码层。

原始的GPT使用pooler函数来获得最终的输出。我们在所有transformer层之上使用一个额外的查询层，通过attention获得最终的embedding。top query层的输入被替换为位置n+1的query embedding。最后的输出再乘以词嵌入矩阵的转置，得到输出概率。对于解码策略，贪心、温度采样、top-k采样、top-p采样和beam search。最后，去标记化将把选中的tokenID变成一个实际的单词。

三、预训练设置

1、代码语料库

训练语料库包含两部分。第一部分来自开源代码数据集，Pile和CodeParrot6。选择了23种流行的编程语言的文件，包括c++，Python,Java,JavaScript,C,Go等。根据每个文件的后缀和它所属的存储库的主要语言来识别其编程语言。第二部分是直接从GitHub公共存储库中抓取的Python、Java和c++的补充数据。

将训练数据分成长度相等的片段。为了帮助模型区分多种语言，我们在每个片段之前以[注释符号]language:[LANG]的形式添加了一个特定于语言的标签，例如#language:Python标记。

2、Tokenization

第一步是将代码片段转换为数值向量。考虑到代码数据中存在大量自然语言注释，变量、函数和类的命名通常是有意义的单词，将代码数据视为文本数据并使用GPT-2 tokenizer，它是一种BPE(Byte Pair Encoding)分词器，使用固定大小的词表和可变长度的字符来处理开放词表问题。初始词汇表大小为50,000，将多个空格编码为额外的标记，以提高编码效率。具体来说，L个空格用<|extratoken_X|>表示，其中X=8+L。由于词汇表包含来自各种自然语言的标记，因此它允许CodeGeeX处理除英语以外的其他语言的标记，如中文。最终的词表大学是v=52,224。经过分词后，任何代码片段或文本描述都可以转换为整数向量。

3、词嵌入和位置嵌入

给定token，下一步是将每个token与单词嵌入相关联。通过在词嵌入矩阵中查找token ID，其中h=5120为隐藏大小，为每个token 获得一个可学习的嵌入。为了获取位置信息，我们还采用了可学习的位置嵌入，将当前位置ID从映射到一个可学习的嵌入，其中nmax=2048为最大序列长度。然后，添加两个嵌入以获得模型的输入嵌入。最后，整个序列可以转化为输入嵌入，其中n是输入序列长度。

四、训练

CodeGeeX在使用 Mindspore (v1.7.0)的 Ascend 910AI处理器 (32GB)集群上进行训练。整个预训练过程在 192个节点和 1536个 AI处理器上花费了两个月的时间，在此期间，模型消耗了 850B个token，相当于5个epochs(213000steps)。为了提高训练效率，采用8路模型并行和192路数据并行，并使用ZeRO-2优化器来进一步降低显存消耗。最终，每个节点上的micro-batch size为16，全局batch size为3072。

使用Adam优化器来优化loss。模型权重采用FP16的格式，为了更高的精度和稳定性layer-norm和softmax使用FP32。模型占用GPU显存为27GB。初始学习率为1e-4，并应用cosine学习率衰减：