PPL.LLM.Kernel.CUDA: 高性能大语言模型CUDA内核库

PPL.LLM.Kernel.CUDA: 加速大语言模型的关键利器

在人工智能和自然语言处理领域,大语言模型(Large Language Models, LLMs)的崛起引发了一场技术革命。然而,这些庞大的模型也带来了巨大的计算挑战。为了应对这一挑战,OpenPPL团队开发了PPL.LLM.Kernel.CUDA - 一个专为大语言模型优化的高性能CUDA内核库。本文将深入探讨这个强大工具的特性、应用场景及其在AI生态系统中的重要地位。

PPL.LLM系统概述

PPL.LLM.Kernel.CUDA是PPL.LLM系统的核心组成部分之一。PPL.LLM是一个全面的大语言模型加速系统,旨在提供从模型导出到高效推理的端到端解决方案。系统的整体架构如下图所示:

PPL.LLM系统架构

从图中可以看出,PPL.LLM.Kernel.CUDA位于整个系统的底层,为上层的模型推理和优化提供了高效的CUDA内核实现。这种分层设计使得系统具有良好的模块化和可扩展性,能够适应不同模型和硬件平台的需求。

核心特性

PPL.LLM.Kernel.CUDA提供了一系列针对大语言模型优化的先进特性:

Flash Attention: 这是一种高效的注意力机制实现,能显著减少内存使用和计算时间。
Split-k Attention: 类似于Flash Decoding,这种技术可以进一步优化解码过程。
Group-query Attention: 通过对查询进行分组,提高了注意力计算的效率。
动态批处理(Dynamic Batching): 又称为Continuous Batching或In-flight Batching,这种技术能够更灵活地处理不同长度的输入序列,提高GPU利用率。
张量并行(Tensor Parallelism): 支持模型的并行计算,充分利用多GPU资源。
图优化(Graph Optimization): 通过对计算图进行优化,提高整体推理效率。
INT8量化: 支持KV缓存和权重的INT8量化,在保持数值精度的同时显著减少内存占用和提高推理速度。

这些特性的组合使得PPL.LLM.Kernel.CUDA能够在保证模型性能的同时,大幅提升推理速度和资源利用效率。

支持的模型

PPL.LLM.Kernel.CUDA目前支持多个主流的大语言模型,包括:

LLaMA 1/2/3系列
ChatGLM 2/3系列
Baichuan-7b
InternLM

这些模型覆盖了学术研究和商业应用的广泛需求,使得PPL.LLM.Kernel.CUDA能够满足不同场景下的应用要求。

性能表现

为了展示PPL.LLM.Kernel.CUDA的强大性能,我们以LLaMA模型为例,在A100 40G GPU上进行了基准测试。测试结果如下表所示:

模型	批大小	张量并行度	输入长度	输出长度	吞吐量 (输出token/秒)
LLaMA 7B	1	1	8	256	85.5
LLaMA 7B	256	1	8	256	8933.3
LLaMA 7B	1	1	1024	1024	82.4
LLaMA 7B	32	1	1024	1024	1247.9
LLaMA 13B	1	2	8	256	76.1
LLaMA 13B	384	2	8	256	8393.0
LLaMA 13B	1	2	1024	1024	73.8
LLaMA 13B	32	2	1024	1024	1199.7
LLaMA 70B	1	8	8	256	42.5
LLaMA 70B	1024	8	8	256	6872.7
LLaMA 70B	1	8	1024	1024	41.3
LLaMA 70B	128	8	1024	1024	1930.7

从这些数据中,我们可以观察到几个关键点:

批处理的重要性: 随着批大小的增加,吞吐量显著提升。例如,对于LLaMA 7B模型,当批大小从1增加到256时,吞吐量提升了超过100倍。
张量并行的效果: 对于更大的模型(如13B和70B),使用张量并行可以有效地分散计算负载,保持较高的吞吐量。
序列长度的影响: 长序列输入会增加计算复杂度,但PPL.LLM.Kernel.CUDA仍能保持较好的性能。例如,LLaMA 70B模型在处理1024长度的输入和输出时,仍能达到每秒1930.7个token的吞吐量。

这些性能数据充分展示了PPL.LLM.Kernel.CUDA在处理大规模语言模型时的高效性和可扩展性。

使用指南

要开始使用PPL.LLM.Kernel.CUDA,您需要先满足以下prerequisites:

Linux操作系统(x86_64或arm64架构)
GCC 9.4.0或更高版本
CMake 3.18或更高版本
Git 2.7.0或更高版本
CUDA Toolkit 11.4或更高版本(推荐11.6)
如果启用NCCL,需要mpich 4.1.2或更高版本

安装步骤如下:

克隆代码仓库:

git clone https://github.com/openppl-public/ppl.llm.kernel.cuda.git

构建项目:

cd ppl.llm.kernel.cuda
./build.sh -DPPLNN_USE_LLM_CUDA=ON -DPPLNN_CUDA_ENABLE_NCCL=ON -DPPLNN_ENABLE_CUDA_JIT=OFF -DPPLNN_CUDA_ARCHITECTURES='"80;86;87"' -DPPLCOMMON_CUDA_ARCHITECTURES='"80;86;87"'

注意:这里的CUDA架构设置(80;86;87)适用于Ampere架构的GPU,您可能需要根据实际硬件进行调整。

导出模型并准备测试数据(以LLaMA为例): 参考ppl.pmx/model_zoo/llama/facebook/README.md中的说明进行操作。
创建benchmark脚本并运行测试: 创建一个名为benchmark.sh的脚本,内容参考项目文档中的示例。然后使用以下命令运行:
```
mpirun -np <MP> benchmark.sh <STEP>
```
其中<MP>是张量并行的大小,<STEP>是测试的步骤。