英文名称: CodeGeeX: A Pre-Trained Model for Code Generation with Multilingual Evaluations on HumanEval-X
中文名称: CodeGeeX：一种用于代码生成的预训练模型，并在HumanEval-X上进行多语言评估
链接: https://arxiv.org/abs/2303.17568
代码: https://github.com/THUDM/CodeGeeX 7.6k Star
     https://github.com/THUDM/CodeGeeX2 6.5k Star
作者: Qinkai Zheng, Xiao Xia, Xu Zou, Yuxiao Dong, Shan Wang, Yufei Xue, Zihan Wang, Lei Shen, Andi Wang, Yang Li, Teng Su, Zhilin Yang, Jie Tang
机构: 清华大学, Zhipu.AI, 华为
日期: 2023-03-30
引用次数: 89

1 读后感

这是一篇比较早的论文，于 2022 年 4 月至 6 月期间，在 1,536 个 Ascend 910 AI 处理器集群上，对 23 种编程语言进行了训练，总共使用了超过 850B 个 token，模型于 2022 年 9 月公开发布。

该论文介绍了 CodeGeeX 的第一代技术，而 CodeGeeX2 则基于 ChatGLM2，底层结构转向了 llama2。新版本没有附带新的论文和开源代码。

通过阅读这篇论文，我们可以了解到开发代码生成的起始阶段，包括：现有对基础架构的调整、数据组织和模型评估的过程。从代码实现的角度来看，CodeGeeX 是完全开源的。它使用 Mindspore 作为深度学习框架，而非主流的 torch。不过，它提供了类似于 torch 的调用方法，所以只需简单浏览即可。

与 copilot 相比，CodeGeeX 免费使用的。我的使用体验是：功能差异不大，只是质量稍逊色一些。

2 摘要

目标：介绍 CodeGeeX，一个拥有 130 亿参数的，支持多编程语言的代码生成模型。

方法：CodeGeeX 用 850 B 个 token 进行了模型预训练，涵盖了 23 种编程语言。

结论：实验结果表明，CodeGeeX 在性能上优于规模相似的多语言代码模型。

3 引言

主要贡献

• 推出了 CodeGeeX，一个 13B 参数的 23 语言代码生成模型，其在代码生成和翻译上超过了同等规模的多语言基线。
• 为 VS Code、JebBrains 和 Tencent Cloud Studio 开发了 CodeGeeX 扩展插件，比 Copilot 有更多功能，包括代码完成、生成、翻译和解释，且能提高 83.4% 用户的编码效率。
• 创建了 HumanEval-X 基准测试，以评估多语言代码模型在代码生成和翻译任务的功能正确性，推动对预训练代码模型的理解和发展。

4 实现

4.1 模型架构

基于 Transformer 框架，采用了仅解码器网络进行自回归语言建模。CodeGeeX 的核心架构是一个 39 层的转换器解码器。在每个 Transformer 层中，应用了多头自注意力机制，然后是 MLP 层、层归一化和残差连接。并使用了 FastGELU 激活函数。

采用 GPT 范式，在大量未标记的代码数据上训练模型。

在顶部查询层和解码方面，原始的 GPT 模型使用池函数来获取最终输出。而文中模型在所有其他转换器层之上使用了一个额外的查询层，并通过注意力来获得最终的嵌入。

4.2 预训练

4.2.1 训练数据

训练语料库包含两部分：

第一部分来自开源代码数据集，包括 Pile（Gao 等人，2020 年）和 CodeParrot3。Pile 包含了 GitHub 上星级超过 100 的公共仓库的子集，从中选择了 23 种流行编程语言的文件，包括 C++，Python，Java，JavaScript，C，Go 等。根据每个文件的后缀和它所属仓库的主要语言来确定编程语言。CodeParrot 是来自 BigQuery 的另一个公开的 Python 数据集。

第二部分是直接从 GitHub 公共仓库中抓取的 Python，Java 和 C++ 的补充数据，这些数据在第一部分中没有出现。选择的仓库至少有一个星级，总大小在 10MB 以内，然后过滤掉那些：1）平均每行超过 100 个字符的文件，2）自动生成的文件，3）字母比例小于 40% 的文件，4）大于 100KB 或小于 1KB 的文件。按照 PEP8 标准格式化 Python 代码。

4.2.2 Tokenization

考虑到 1）代码数据中存在大量自然语言注释，2）变量、函数和类的命名通常是有意义的词，将代码数据与文本数据相同，并应用 GPT-2 标记器。由于词汇表包含来自各种自然语言的标记，因此它允许 CodeGeeX 处理英语以外的语言的标记，如中文、法语、俄语、日语等。最终的词汇量为 52,224 。

4.2.3 输入词和位置嵌入

在给定 tokens 后，每个 token 会与一个词嵌入相关联，同时也会获取位置信息的嵌入。这两种嵌入相加后形成模型的输入嵌入，最后整个序列被转换为输入嵌入。

4.3 训练

在 Ascend 910 上进行并行训练。CodeGeeX 在带有 Mindspore（v1.7.0）的 Ascend 910 AI 处理器（32GB）集群上进行了训练。与 NVIDIA GPU 和 PyTorch/TensorFlow 相比，Ascend 和 Mindspore 相对较新。整个预训练过程需要两个月的时间，在 192 个节点和 1,536 个 AI 处理器上，850B 代币，相当于 5+ 个 epoch（213,000 步）。

4.4 快速推理

为了提供预训练的 CodeGeeX，实现了一个纯 PyTorch 版本的 CodeGeeX，它支持在 NVIDIA GPU 上进行推理。为了实现快速且节省内存的推理，将量化和加速技术应用于预训练的 CodeGeeX。

量化前后对比如下：

5 评测

主实验结果如下：

后面的 CodeGeeX2 相比 CodeGeeX 又好很多，下图取自 CodeGeeX2 github。