英文名称: 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。