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，感兴趣的小伙伴可以留言，后续我再推出更详细的解读。

• CodeGeeX 代码链接：https://github.com/THUDM/CodeGeeX （点击右下角 BBuf 的头像就可以找到 oneflow 的 pr）

• OneFlow 代码链接：https://github.com/Oneflow-Inc/oneflow

0x1. 优化后的结果

我们在 A100 PCIE-40G 上对比了分别使用 PyTorch，FasterTransformer 以及 Oneflow 推理 CodeGeeX 模型的耗时情况，FP16 模式推理速度结果如下：

INT8 模式的推理速度如下：

可以看到无论是在 FP16 模式还是 INT8 模式，OneFlow 均取得了最好的性能结果。也许有些读者会提出似一个疑问，似乎 OneFlow 的性能并没有超越 FasterTransformer 太多，选择 OneFlow 的好处是？我个人认为由于 C++以及手动插入集合通信的原因 FasterTransformer 的适配难度是相对比较大的，特别是多卡模式，而 OneFlow 不仅拥有和 PyTorch 一模一样的用户体验并且扩展到多卡时不需要用户手动管理集合通信的问题，用户体验拉满。

除了性能优势，OneFlow 也可以节省一些显存资源消耗，详细的信息可以点击这个链接查看：https://github.com/THUDM/CodeGeeX/pull/87 。

0x2. 优化手段解析

针对 CodeGeeX 大模型的推理，OneFlow 做了什么优化可以超越 NVIDIA FasterTransformer 库的推理速度呢？

• quick_gelu 融合优化。https://github.com/THUDM/CodeGeeX/blob/main/codegeex/oneflow/codegeex_model.py#L7-L11 指的是将x / (1 + torch.exp(-1.702 * torch.abs(x))) * torch.exp(0.851 * (x - torch.abs(x))) 这个 elementwise 操作组合成的 pattern 融合成一个算子，在 oneflow 中为flow._C.quick_gelu。

• grouped_matmul_bias 优化。https://github.com/THUDM/CodeGeeX/blob/main/codegeex/oneflow/codegeex_model.py#L101-L108 指的是将一堆同时执行并且数据没有前后依赖关系的 matmul+bias_add 算子融合成一个 cuda kernel，降低 kernel launch 的开销。https://github.com/Oneflow-Inc/oneflow/pull/9413。

• 更高效的 fused attention kernel（在 oneflow 中使用flow._C.fused_multi_head_attention_inference_v2调用）。在 oneflow 中引入了 cutlass 的 fmha 以及 TensorRT 的 FlashAttention 实现，可以在不同的数据规模调用最优的 fmha 实现。在此基础上 oneflow 针对 Q，K，V 可能存在的不同数据排布进行优化，具体来说 oneflow 的 fused_multi_head_attention_inference_v2 接口支持手动配置 Q，K，V 这三个输入 tensor 的数据排布。比如在 CodeGeeX 里面，Q，K，V 的 shape 是[seq_lenght, batch_size, num_heads * hidden_size_per_attention_head]，我们就可以直接把 Q，K，V 的数据排布配置成MB(HK)，并且输出的数据排布也配置成 MB(HK)，这样就可以避免在把 Q，K，V 传入 fused_multi_head_attention_inference_v2 之前需要额外做的 reshape 带来的开销了，同样输出 Tensor 的 reshape 开销也可以避免。https://github.com/THUDM/CodeGeeX/blob/main/codegeex/oneflow/codegeex_model.py#L253-L264 。这部分的 cuda 实现分成很多 pr，这里指一下路：https://github.com/Oneflow-Inc/oneflow/pull/9950 & https://github.com/Oneflow-Inc/oneflow/pull/9933。

• CodeGeeX 和大多数的自回归模型一样有一个增量推理阶段，需要把当前的 key,value 和上一轮的 key,value concat 起来，也就是：https://github.com/THUDM/CodeGeeX/blob/main/codegeex/oneflow/codegeex_model.py#L135-L140 。针对这个特殊的操作，我们也开发了一个可以配置输入输出数据排布的 fuse kernel，把两个 concat 操作融合起来降低 kernel launch 以及 reshape 的开销。https://github.com/THUDM/CodeGeeX/blob/main/codegeex/oneflow/codegeex_model.py#L239 。在 oneflow 中对应 https://github.com/Oneflow-Inc/oneflow/pull/9963 。

• fused matmul+bias。https://github.com/THUDM/CodeGeeX/blob/main/tests/test_inference_oneflow.py#L14 。具体来说就是将 Linear 中的 matmul 和 bias_add 融合在一起。https://github.com/Oneflow-Inc/oneflow/pull/9369。

上述优化既适用于 FP16 模式，也适用于 INT8 模式，接下来我们聊一下 INT8 weight only quantization 的 motivation 以及优化。经过调研，FasterTransformer 的 INT8 模式采用了 weight only quantization 的方式，也就是只对 Linear 层的权重进行量化，但是在计算的时候仍然要反量化回 FP16 和 Activation 进行矩阵乘计算。按道理来说，加入了反量化之后速度应该变慢才对，为什么这里使用了 INT8 weight quantization 之后反而能加速最终的推理速度呢？这是因为在这个网络中，推理时的 batch_size 以及 seq_length 都是 1，这个时候的矩阵乘法退化到了一个向量和一个矩阵相乘的情况，实际上类似于卷积神经网络中的全连接层，是一个典型的访存密集型算子。所以这里对 weight 进行反量化和矩阵乘法可以 fuse 到一起来进行加速（原因是减少了访存）。在 oneflow 中的实现对应：https://github.com/Oneflow-Inc/oneflow/pull/9900 。然后我基于这个算子在 CodeGeeX 中实现了 OneFlow INT8 版本的推理脚本：https://github.com/THUDM/CodeGeeX/blob/main/codegeex/quantization/quantize_oneflow.py

0x3. 总结

​至此，我分享完了我们团队最近加速 CodeGeeX 百亿参数大模型推理的所有优化技巧，相信对要做 LLM 大模型的推理的小伙伴会有帮助。本着开源精神，请给 oneflow 点击 star 再研究相关优化。此外，更多的优化解读我也会放到个人仓库：https://github.com/BBuf/how-to-optim-algorithm-in-cuda ，欢迎大家关注。

本文源自：“ GiantPandaCV”公众号