本文主要总结两篇文章:Deja Vu 和 Apple 的 LLM in a flash。这两篇文章的内容都是端侧推理加速的尝试,他们主要使用了大致的思路–利用MLP的稀疏性,各自的工程实现各有一些创新。

端侧推理有着比较大的应用前景,随着端侧设备的算力跟进,端侧设备已经具有了运行7B模型的能力。在端侧运行小参数模型可以极大的减少云端的压力,从而减少运营成本。相比于云端的大模型,端侧大模型处理复杂问题能力不足,所以端侧和云端应该是相辅相成的。轻量级任务给端侧,需要长逻辑理解的任务交给云端。

端侧和云的协同工作,也是一个很好的研究方向。

0x01 Deja Vu

1. 问题分析和动机

作者通过分析OPT-175B模型的上下文稀疏性发现对于大部分的Transformer Layer,他们的稀疏性都在85%左右。上下文稀疏性就是:对于特定的输入,仅有一小部分的模型参数对最终结果有着重要的影响。 如图1-3所示:

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time
Fig 1. Contextual Sparsity

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time
Fig 3. Contextual sparsity in Attention Head

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time
Fig 3. Contextual sparsity in MLP Block

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time

在论文中,作者通过以下的步骤来计算上下文稀疏性:

  1. 首先,对模型进行一次前向传播,记录下在计算过程中哪些注意力头和MLP神经元产生了较大的输出范数(output norms)。
  2. 然后,对每个输入样本进行第二次前向传播,但这次只使用第一次传播中记录的参数子集来计算输出。
  3. 通过比较两次前向传播的结果,作者们验证了即使只使用一部分参数,模型的输出也能保持与完整模型相当的准确性。这意味着对于特定的输入,模型展现出了上下文稀疏性。

相比较于静态分析的剪枝方法,基于上下文稀疏性来进行剪枝可以获得更好的Acc:

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time
Fig 4. Accuracy-Efficiency Trade-offs

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time

上下文稀疏化根据当前输入动态调整激活的参数,实现了对模型的精准剪枝。这种适应性的方法比固定不变的剪枝策略更加灵活高效,因为它能够智能地选择对特定输入最为关键的参数,从而在牺牲极少模型性能的前提下,有效提升推理速度和计算效率。

作者在实验的过程中还发现:Attention Head的每一个Head的输出对的分布是不同的,有的Head中Token对应的Attention Score很高,有的则很小。 作者假设,注意力机制是在做mean shift聚类。每一次过Attention Block就是在完成一次聚类迭代。作者将Attention Score打印了出来,见图5。

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time
Fig 5. visualize the attention scores of three different heads for an exemplary sentence.

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time

作者们还观察到了层与层之间数据变化的特性。他们发现,在激活层中,尽管数据会经历线性和非线性变换,但相邻层之间的嵌入向量或激活向量仍保持高度的余弦相似性。 以OPT-175B模型为例,除了第一层外,任意两个连续层之间的相似度接近0.99,这表明在模型的传递过程中,输入的嵌入表示在层间的改变是渐进的,且在更深层的网络中这种特性更为显著。

此外,作者们还探讨了导致这种高相似性的原因,指出这是由于Transformer中的残差连接造成的。残差连接允许信息在层与层之间直接传递,减少了因层内复杂变换导致的信息丢失。作者们还提供了数学上的分析,比如证明了在特定条件下,模型的层间变换具有收缩性质,这有助于维持信息的稳定性,从而使得层间的嵌入向量保持高度相似。这些发现对于理解和改进深度学习模型的效率和性能具有重要意义。

2. 有效神经元预测算法

作者使用了一个较小的额外的神经网络分类器来预测每个层的上下文稀疏性。这个分类器预测每个Attention Head和MLP的参数在给定输入下的活跃度,然后做动态的剪枝。

Deja Vu在预测的时候使用的是异步的策略,让分类器和模型的计算并行执行。之所以可以并行执行,是因为“在激活层中,尽管数据会经历线性和非线性变换,但相邻层之间的嵌入向量或激活向量仍保持高度的余弦相似性。 ” 这个重要结。因为在做计算的时候,是tiling的,所以我们可以通过活跃度选择性的加载我们需要的tile到内存里面去,从而减少内存的IO。

本文的结果也很惊艳:

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time
Fig 6. Accuracy Trend for DEJAVU-OPT-175B.

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time
Fig 7. Average per-token latency (ms).

Deja Vu: Contextual Sparsity for Efficient LLMs at Inference Time

sparsity percentage在70%左右的时候Acc些许下降,相比于FasterTransformer,有一倍latency下降。

0x00 LLM in a Flash

一个7B半精度LLM,放在DRAM中需要14GB。纵使4bit量化,也要2.5GB+的DRAM,这对于端侧设备来说是一个不小的负担。这篇文章就沿用了Deja Vu的思路,并在工程上做了很多有意思的工作。本文的动机其实很常见,那就是:DRAM放不下,那就从Flash中取就行了,类似于SWAP?但是对于LLM的任务,对于DRAM带宽的要求很高,更不用说Flash Memory了。文中给出了各个部件的带宽,见图8-9。

LLM in a flash: Efficient Large Language Model Inference with Limited Memory
Fig 8. Bandwidth in a unified memory architecture

LLM in a flash: Efficient Large Language Model Inference with Limited Memory

LLM in a flash: Efficient Large Language Model Inference with Limited Memory
Fig 9. Random read throughput of flash memory

LLM in a flash: Efficient Large Language Model Inference with Limited Memory

这篇文章使用的方法和Deja Vu非常的相似,但是使用的是类似于LoRA的分类器。

LLM in a flash: Efficient Large Language Model Inference with Limited Memory
Fig 10. Low Rank Predictor

LLM in a flash: Efficient Large Language Model Inference with Limited Memory

本文主要的难点是:

  1. 如何快速准确的找出哪些模型的参数是重要的
  2. 如何减少Flash Memory的负载压力,或者说如何减少Flash Memory和DRAM之间的数据传输。实际上第一点也是减少数据传输的手段。

1. Selective Persistence Strategy

LLM的参数主要是两个部分,Attention部分的参数和MLP部分的参数,其中MLP部分的参数占比非常的大。Apple在这里的方案是将Attention和Embedding层常驻在内存中,MLP部分的参数使用LoRA预测器(应该是灵感来源于Deja Vu)来进行预测。

2. Sliding Window

每一次LLM前向推理的时候,都需要过LoRA,加载需要的参数,但是这个参数加载的过程是随机读取的,对于Flash来说并不友好。比如第一个token在前向的第一个MLP需要{0, 1, 7, 9}位置的参数,第二个token在前向推理的时候要{2, 3, 7, 9}位置的参数。那么我们其实可以keep参数{7, 9}来避免二次读取。这就是LRU的思想,Apple的方案在这里借鉴了LRU的方法。Sliding Window的核心思想就是LRU:保留处理过去k个token时的激活神经元所对应的参数在DRAM中,并在处理当前token时只对部分多余的参数进行删除和对缺少的参数进行加载。

LLM in a flash: Efficient Large Language Model Inference with Limited Memory
Fig 11. Sliding window

LLM in a flash: Efficient Large Language Model Inference with Limited Memory