2022年底,OpenAI发布了跨时代的ChatGPT应用。它的成功使大模型成为AI发展的主旋律,在极短的时间内改变了AI产业的格局。随着GPT-4、Gemini、Sora、Claude3、Kimi等一系列大模型的陆续发布,大模型能力迅速提升,甚至更为强大的通用人工智能(artificial general intelligence,AGI)已初见端倪。
大模型以其强大的理解和生成能力正在深刻改变我们对人工智能的认知和应用,但其高昂的推理成本也阻碍了技术落地。因此,优化大模型的推理性能成为业界研究的热点。本文中,我们试图对ChatGPT发布以来大模型推理优化关键技术做出综述,厘清技术全貌及发展态势,以便读者做出更好的判断和预测。
大模型推理性能优化主要以提高吞吐量和降低时延为目的,关键技术可以划分为:内存管理、算子融合、模型压缩、并行推理、服务调度优化及新兴技术。
内存管理
KV Cache是大模型推理性能优化最常用的技术。该技术在不影响任何计算精度的前提下,通过空间换时间,大幅提升推理性能。Transformer解码器使用自回归产生输出,即每次推理只会预测输出一个token,执行多次后完成全部输出。前后两次的输入只相差一个token,这就存在大量重复计算。KV Cache技术将每个token可复用的K和Q向量结果保存下来复用,将计算复杂度从O(n2)降低为O(n)。
Paged Attention技术将操作系统中的分页内存管理应用到KV Cache的管理中,节约了60%~80%的显存,从而支持更大的batch-size,将吞吐率提升了22倍。具体来讲,Paged Attention首先将每个序列的KV Cache分成若干块,每个块包含固定数量token的键和值,然后计算出当前软硬件环境下KV Cache可用的最大空间,并预先申请缓存空间。在推理过程中,通过维护一个逻辑块到物理块的映射表,使多个逻辑块对应一个物理块,并使用引用计数标记物理块被引用的次数,从而实现将地址不连续的物理块串联在一起统一管理。
算子融合
算子融合是深度学习模型推理的一种典型优化技术,旨 在 通 过 减 少 计 算 过 程 中 的 访 存 次 数 和 统 一 计 算 架 构 (CUDA)Kernel的启动耗时,达到提升模型推理性能的目的(见图1)。针对Transformer的结构特点,算子融合主要分为4类:归一化层和QKV横向融合,自注意力计算融合,残差连接、归一化层、全连接层和激活层融合,偏置加法和残差连接融合。
中兴通讯在vLLM上实现了针对多查询注意力结构的QKV通用矩阵乘法(GEMM)横向算子融合,以及多层感知机(MLP)中的全连接层(FC)+激活融合,性能明显提升,见表1和表2。上述算法的相关代码实现已合入vLLM社区。
模型压缩
模型压缩技术是指在不影响模型精度的情况下,通过缩小模型规模和计算量来提高模型的运行效率,其中模型量化是最具实用性的技术。
SmoothQuant是典型的8bit LLM量化方法,该方法引入了逐通道缩放变换,有效地平滑了幅度,这使得模型更易于量化。激活感知权重量化(AWQ)和生成式预训练Transformer GPTQ是典型的权重量化方法,且权重量化的是group粒度。GPTQ提出了一种基于近似二阶信息的新型分层量化技术,使得每个权重的比特宽度减少到3或4位。AWQ的研究人员发现,对于LLM的性能,权重并不是同等重要的,仅保护1%的显著权重可以大大减少量化误差。
中兴通讯提出了SmoothQuant+4bit权重量化训练后量化(PTQ)算法。不同于AWQ对单个层搜索量化参数,SmoothQuant+对整个模型搜索量化参数,并对整个模型进行同一个参数平滑激活,这样能够从模型整体减少量化误差,且搜索效率更高。SmoothQuant+在LLaMA系列模型可以得到比AWQ更好的精度(见表3),同时在性能上也优于AWQ,对应的推理核已开源。
随着大模型上下文长度的增加,KV Cache占用的显存将超过权重和激活,因此对KV Cache进行量化可以显著降低大模型在长上下文推理时的资源占用,从而允许系统支撑更多的并发请求数和吞吐率。
并行推理
当大模型参数量超过单一计算设备所能容纳的上限时,则需要使用分布式并行推理技术。并行推理可以使用模型并行和流水线并行,而模型并行由于可节省显存资源、可降低单用户时延等优势,成为首选的并行方式。业界最流行的模型并行方案来自Megatron-LM,它的开发者针对Self-Attention和MLP分别设计了简洁高效的模型并行方案。而节点间带宽对模型并行效率有较大影响,高速串行计算机扩展总线标准(PCIE)的理论带宽为32~64Gbps,通常可以满足大模型并行推理需求。
服务调度优化
服务调度优化主要考虑的是系统同时为多个用户服务时如何尽可能地提升资源利用率。Continuous Batching和Dynamic Batching主要围绕提高可并发的Batchsize来提高吞吐量,异步Tokenize/Detokenize则通过多线程方式将Tokenize/Detokenize执行与模型推理过程时间交叠,从而实现降低时延目的。
Continuous Batching可以将传统batch粒度的任务调度细化为step级别的调度,这解决了不同长短序列无法合并到同一个batch的问题,大幅提升推理效率和用户体验,目前已在HuggingFace TGI、vLLM、TensorRT-LLM等多个推理框架中实现。
新兴技术
我们将传统优化技术引入大模型推理的同时,也在探索从大模型自回归解码特点出发,通过调整推理执行过程来进一步提升推理性能。并行推测解码作为新兴的推理技术,可以在不损失精度的前提下提高推理速度。
投机采样是一种并行推测解码算法,开创了“小成本生成+大模型验证”的推理技术路线。该算法在已有大模型的基础上,引入一个小模型执行串行解码来提升速度,原大模型执行并行评估采样,保证生成质量,这在保证精度一致性的同时降低了大模型解码的次数,进而提升了推理效率。由于投机采样算法的巨大潜力,有多项工作在其基础上研究改进。但投机采样的推理方式并不适用于所有的应用场景。例如,在文学艺术类的诗词等应用场景,大小模型生成的结果概率分布相差较大;对于代码生成的场景,投机采样比较适合。随着业界研究的深入,投机采样会成为大语言模型推理的必备优化技术。
中兴通讯研发了星云编程大模型,通过上述技术优化,实现显存节省70%,单GPU卡吞吐量提升3倍,推理时延降低一半,推理成本降低75%左右。目前中兴通讯内部已经建设研发大模型推理集群,将大模型辅助编程集成到研发IDE环境中,每日超过1.3万员工使用编程大模型进行开发,日生成代码超过百万行。
此外,在城市治理、工业、矿山、交通等行业存在多种大模型应用场景,除数据中心部署外,在网络边缘、现场设备上也存在部署大模型需求。中兴通迅利用量化压缩技术降低大模型资源占用,降低了大模型部署成本;提供模型编译迁移工具适配不同硬件平台,将大模型部署在边缘一体机和AIBOX上,扩大了大模型应用范围。通过上述技术,星云大模型已经在城市小散工程监管、矿山安全生产、交通管控等实际落地,让行业客户真正用得起大模型。
随着ChatGPT热度的逐渐褪去,对大模型的投资也逐渐趋于理性。大模型如何产生真正的商业价值成为全行业都在思考、探索的问题。随着大模型规模的不断增加,模型能力在提升的同时,算力成本也在不断飙升,这给大模型长期可持续发展带来了不确定性,因此以实现更低成本算力和更高效率算法为目标的核心技术亟待突破。大模型机遇与挑战并存,加速发展的趋势在中长期不会改变。