CodeGeeX 130亿参数大模型的调优笔记：比FasterTransformer更快的解决方案

最新推荐文章于 2025-11-20 01:49:20 发布

原创

最新推荐文章于 2025-11-20 01:49:20 发布 · 1.4k 阅读

3 ·

CC 4.0 BY-SA版权

文章标签：

#深度学习 #python #人工智能

文章介绍了CodeGeeX，一个由清华大学等机构开发的代码生成模型，以及OneFlow如何通过一系列优化手段在FP16和INT8模式下提高其推理速度，超过了FasterTransformer。这些优化包括quick_gelu融合、grouped_matmul_bias、fused_attention_kernel等，并讨论了INT8权重量化的优势。OneFlow的优化不仅提升了性能，还简化了多卡模式的使用体验，并节省了显存资源。

0x0 背景

相信大家都使用或者听说过github copilot这个高效的代码生成工具。CodeGeeX类似于github copilot，是由清华大学，北京智源研究院，智谱AI等机构共同开发的一个拥有130亿参数的多编程语言代码生成预训练模型。它在vscode上也提供了插件，可以直接安装使用，我个人体验了一下代码生成的功能还不错。此外除了代码生成，CodeGeeX还可以做代码加注释，不同语言翻译（比如把c++代码翻译为python）等，感兴趣的读者可以体验一下。并且可以在 https://models.aminer.cn/codegeex/blog/index_zh.html 这个官方博客上查看更多详细信息。

为了说明oneflow在大模型训练和推理上的高效性，继上次对glm10b模型的训练优化工作之后，我们对CodeGeeX模型的推理进行优化。在oneflow团队的优化下，CodeGeeX可以使用oneflow的后端进行推理并且在FP16和INT8模式的推理速度均可以超过CodeGeeX团队基于FasterTransformer的方案（基于NVIDIA A100显卡进行测试）。oneflow的推理方案已经upstream CodeGeeX的主分支，欢迎小伙伴查看。需要指出的是本文用到的大多数cuda优化手段均由oneflow的柳俊丞大佬提供，在此致敬。本着开源精神，本文将展示一下我们的优化结果并且解析一下我们的优化手段，和大家共同探讨学习。介于篇幅原因，在解析优化手段时，我们会简单介绍一下优化的原理并给出代码链接。但不会详细阅读优化涉及到的cuda kernel，感兴趣的小伙伴可以留言，后续我再推出更详细的解读。