CPU推理提升4到5倍,苹果用闪存加速大模型推理,Siri 2.0要来了?近年来,GPT-3、OPT 和 PaLM 等大型语言模型(LLM)在广泛的 NLP 任务中表现出了强大的性能。不...

贴文 2个月前
3.2K 21

CPU推理提升4到5倍,苹果用闪存加速大模型推理,Siri 2.0要来了?

近年来,GPT-3、OPT 和 PaLM 等大型语言模型(LLM)在广泛的 NLP 任务中表现出了强大的性能。不过,这些能力伴随着大量计算和内存推理需求,毕竟大型语言模型可能包含数千亿甚至万亿参数,使得高效加载和运行变得有挑战性,尤其是在资源有限的设备上。

当前标准的应对方案是将整个模型加载到 DRAM 中进行推理,然而这种做法严重限制了可以运行的最大模型尺寸。举个例子,70 亿参数的模型需要 14GB 以上的内存才能加载半精度浮点格式的参数,这超出了大多数边缘设备的能力。

为了解决这种局限性,苹果的研究者提出在闪存中存储模型参数,至少比 DRAM 大了一个数量级。接着在推理中,他们直接并巧妙地从闪存加载所需参数,不再需要将整个模型拟合到 DRAM 中。

这种方法基于最近的工作构建,这些工作表明 LLM 在前馈网络(FFN)层中表现出高度稀疏性,其中 OPT、Falcon 等模型的稀疏性更是超过 90%。因此,研究者利用这种稀疏性, 有选择地仅从闪存中加载具有非零输入或预测具有非零输出的参数。

论文地址:网页链接

具体来讲,研究者讨论了一种受硬件启发的成本模型,其中包括闪存、DRAM 和计算核心(CPU 或 GPU)。接着引入两种互补技术来最小化数据传输、最大化闪存吞吐量:

窗口:仅为前几个 token 加载参数,并重用最近计算的 tokens 的激活。这种滑动窗口方法减少了加载权重的 IO 请求数量;行列捆绑:存储上投影和下投影层的串联行和列,以读取闪存的更大连续块。这将通过读取更大的块来增加吞吐量。

为了进一步最小化从闪存传输到 DRAM 的权重数量,研究者还设法预测 FFN 稀疏性并避免加载归零参数。结合使用窗口和稀疏性预测可以为每个推理查询仅加载 2% 的闪存 FFN 层。他们还提出了静态内存预分配,最大限度减少了 DRAM 内的传输并减少了推理延迟。

本文的闪存加载成本模型在加载更好数据与读取更大块之间取得了平衡。与 CPU 和 GPU 中的 naive 实现相比,优化该成本模型并有选择地按需加载参数的闪存策略可以运行两倍于 DRAM 容量的模型,并将推理速度分别提升 4-5 倍和 20-25 倍。

有人评价称,这项工作会让 iOS 开发更加有趣 网页链接

21 条评论