#大语言模型系统侧优化简述
1. 前言:
最近有很长的时间在看推理框架的源码,多数都是细节,这里算是自己简单总结下总体相关的内容。
受限于没有集群,所涉及的部分可能存在一定偏差,如有问题请指正。
这里只是个人观点,仅供参考。
Towards Efficient Generative Large Language Model Serving: A Survey from Algorithms to Systems
还是基于这篇论文来做一个简要总结。
由于能力所限,这里只讨论:
System Optimizations部分的内容。
2. 量化
由于LLM的训练数据的规模和成本相关的考虑,目前提供的量化方案都是训练后量化,而不是感知量化。
训练后量化目前最主流的两个算法是:GTPQ和AWQ。
知乎有大量的文章介绍量化算法。这里不做赘述。
其优势在于:
- 大幅降低内存占用,尤其是WEIGHTS的内存占用,对于小模型+量化在边缘端优势明显。
- 可以大幅减少模型IO,实际提升推理速度。
HUGGINGFACE这些都有集成,算是非常成熟的解决方案。
3. 算子优化
目前算子优化主要集中在两个层面:
- ATTN部分的KVCACHE
- GEMM部分。
其中ATTN的优化主流以FLASH-ATTN和FLASH-INFER两个开源库为底层。目前FLASH-ATTN-2底层使用CUTE进行代码编写,主要适配于FP16的数据处理。
而在FP8的数据结构下,目前FLASH-INFER的性能有一定优势。(没实际测试过。。。VLLM有测试过)
对于GEMM部分(不限于MLP),目前了解下来,则是以MARLIN+量化为标准的FP16*INT4计算库为最卷。
说点个人的想法:
- 个人感受到的趋势是:CUTE很可能会是未来算子的范式。
举个例子,marlin使用纯PTX的指令进行编写,其代码可能性尤其是地址相关的可读性是真的很逆天。
同样flash-attn的代码,虽然CUTE也需要入门成本,但是代码的复杂度还是降低了非常的多。
- 算子层的调度也可以很复杂
如marlin的多级规约,很有想法。
4. 并行
目前vLLM内部的TP,DP并行已经相对完善。
PP并行目前还不太行。
目前TP的并行基本已经范式。简单可以看ATTN部分的这两张图:
分别支持在算子层面按Row,Col方向就行切片。后面再做REDUCE。
PP和TP不同,PP是直接按LAYER进行切片。
TP需要将额外的激活层的数据交给下一个节点进行处理,这是增加了开销,但是这部分数据较少。
因此这里会有一个经验之说,如果集群内互联使用NVLINK的高速带宽进行传递,就用DP并行,如裹网络条件较差,则建议使用PP并行。
dp并行现在主流框架的做的都挺好,PP并行推荐Megatron-LM。
5. 调度策略
个人认为这会是未来一段时间内框架最卷(变化最快)的部分。
LLM有一个核心问题,PREFILL是计算密集型,而DECODE是访存密集型。而DECODE又依赖PREFILL,同时DECODE的长度又是变长的。这就注定了不太容易有一套简单粗暴且容易的方法能兼顾两者。
目前vLLM主流调度还是PREFILL FIRST的逻辑。
PREFILL/DECODE分离式目前已经在进行的。
chunked-prefills实际效果需要时间来检验。
到chunked-prefills这步,个人认为算是深水区了。
6. 工程调度
vllm在0.6.0版本加入了mutil-step,0.6.2加入了output aysnc stream。
作为菜鸡的我认为卷的应该差不多了。
快一统了。。。
7. 内存相关
prefix caching
长prompt优化。牛。
8. 编译
其实这部分最近接触的很少。
我以前做CV,编译的内容还是有一些的。
靠制定各种规则来做FUSE-OPS的操作在LLM框架里很少见到。--也有可能是我孤陋寡闻。
目前个人认知里,TENSORRT-LLM推理性能上并没有明显优势,不像TENSORRT在CV任务里那么明显了。
9. 其他优化
spec decoding
真正大模型才用的上优化,且可能存在带宽瓶颈。但确实是个很好的优化方向。
如果说缺点的化,这个优化和mutil-step有很大的冲突。同时兼容的代价有些大。
spec decoding需要再观察观察。