IPADS工作中SOSP，LLM推理的新探索？

iPad

今日发觉之前读过的一些paper被SOSP录用了，挑几篇简单说说我的浅见，若有问题还请大家指出。iPadS的这篇工作大家应该都比较熟悉了，这次终于被SOSP录用，那我就贴些之前的笔记来讲一讲。简而言之，将稀疏性与CPU和GPU混合推理的潜力相结合。LLM Sparsity的问题之前Deja Vu也提到过，这篇文章也是沿着这个方向继续探索应用场景。但有一个局限性，这两篇文章似乎主要研究的是以ReLU/ReGLU为激活函数的FFN层（实验设置采用的是OPT、Falcon - ReLU和LLaMA - ReGLU）。可以想见，对于GeLU模型而言，肯定要重新进行一遍微调才能适配。这篇文章的动机是在消费级个人电脑（consumer - grade PC）上进行大型语言模型（LLM）推理，也就是显存（VRAM）不足以容纳整个模型的情况。此前也有一些通过卸载（offloading）来减少显存消耗的工作，笔者比较了解llama.cpp卸载部分层（layer）并进行混合推理的方式，不过这种方式的局限性非常明显：在显存极度受限的情况下，将大部分层都交给CPU处理会极大地降低推理速度。为解决这些问题，这一工作利用上文提及的大型语言模型稀疏性（LLM Sparsity）对卸载进行更精细的调度，这样就能在GPU仅存储少量权重时仍然保持高效。

具体的调度策略包含两个步骤。第一步，运用offline的profiler，确定哪些神经元的影响最大，这被形式化为一个ILP。第二步，利用online的predictor预测哪些神经元会被激活，然后预加载到GPU。这里的思路和Deja Vu较为相似，就不再详细说明了。

在Implementation（实现）和Experiment（实验）部分，此工作以llama.cpp为基础。或许是由于llama.cpp本身具备offload GPU（卸载GPU）功能，并且能够直接当作baseline（基准）来进行对比。

电脑

Evaluation部分还有一些值得关注之处，像prefill的speedup较小（在computation - intensive时CPU会成为明显的瓶颈）等情况，在此就不再详细叙述了。总体而言，这篇工作聚焦于低显存（low VRAM）场景，展现了混合推理与稀疏性（sparsity）的潜力。我个人觉得，这两点后续都有诸多值得探索的应用场景。此外，还可以更多地去考量稀疏性与在显存受限（memory constrAInt）时更常用的量化（quantization）方法之间的比较或者结合（在评估（evaluation）中已经进行了简单的int4量化结合实验）。续作PowerInfer - 2所侧重的场景更倾向于手机这种具有统一内存以及相对特殊的异构场景。不过我还没有仔细研究，之后可能会补充这方面的内容。这篇工作又是金老师组做的，好像是今年第四篇OSDI/SOSP了，真的很厉害！集群场景中sequence parallelism的动态调度和通信优化（简而言之）。

细节上，在scale down（缩减规模）这部分，因为sequence parallel（序列并行）在prefill（预填充）计算时原本就需要传输key - value tensor（键值张量），所以当并行度降低时，可以直接预先把一些为计算而传输过来的key - value tensor存储在本地，防止后续再次通信产生额外开销。

scale up部分把Flash - Decoding思想扩展到多卡情境。当要提升并行度时，只需指定一个有空闲显存的GPU加入并行组，并让它成为后续计算的主节点（master）：（1）先在主节点进行QKV投影计算，得出注意力（attention）的输入张量；（2）把查询张量（query tensor）发送到其他存有该请求上文KV缓存的GPU进行快速解码（flash decoding）；（3）再把计算结果发回主节点以进行后续计算（如前馈神经网络FFN等）。因为每个工作节点/ GPU能同时作为一个请求的主节点和另一个请求的从节点，所以文中将这种方法称为多主分布式解码（multi - master distributed decoding）。

下面是调度方面的内容，这篇工作的调度分四步进行。第一步，在当前迭代（iteration）中选择开始预填充（prefill）的请求。选择时的考虑因素有：一是要防止键值（KV）缓存爆显存；二是要尽可能让GPU计算达到满载；三是要考量是否会抢占当前正在解码（decode）的请求以及抢占的代价。第二步，为选出来的预填充请求分配合适的弹性实例（elastic instance），这里主要考虑迁移（migration）开销。第三步，将这些预填充请求和实例进行映射。采用相对启发式的方式来进行前面提到的弹性扩缩（elastic scale up/down）。其实这一段理论分析有部分我没怎么弄明白，四边形不等式单调性证明那部分尤其如此，等哪天有空了重新看看OI笔记看能否搞懂。

在最后的实验环节，将其与vLLM、SplitFuse和DistServe作对比，能够发现这一工作采用的更具动态性的调度策略，在负载较高时，相较于其他基准（baseline）有十分显著的提升。此外，还在单机多卡和跨节点的场景下开展了实验，再结合后续关于规模（scale）的消融实验（ablation），共同表明文中提及的优化（optimization）能够很好地掩盖通信开销。

还有些工作文件没发布，之后有空会继续更新。