WD：Flash加速AI推理在端侧落地

摄影彬哥哥 2024-10-12 10:35 发布于中国

关键要点

使用闪存加速边缘设备上大型语言模型（LLM）采用。
在客户端设备上运行LLM时，需要解决内存墙问题。
通过将部分LLM加载到GPU VRAM中，可以减少对系统内存的需求。
利用闪存低延迟和高速度，可以实现更高效的参数加载和计算。
通过使用稀疏性预测算法和行列捆绑技术，可以进一步优化LLM的加载和计算效率。

内存墙、算力与带宽

内存墙效应

模型参数 410x / 2yrs
硬件内存 2x / 2yrs

算力、DRAM容量、片上互联速率三者的增幅差距。

算力 3x / 2yrs
内存带宽 1.6x / 2yrs
片上互联带宽 1.4x / 2yrs

想表达什么？

与算力紧缺相比，数据管道读取带宽、通信效率是限制AI训练更关键因素。

端侧设备模型推理挑战

AI应用在端侧设备落地过程遇到的问题

SLM 模型虽已显著压缩，但与当前端侧设备的DRAM容量相比，仍明显超出。

图示 Apple（业内著名内存抠搜）和 Windows 两大OS在融合大模型过程面临挑战。

端侧toC市场对价格非常敏感，提高VRAM以支持客户端推理的方式被认为是不经济的。

下图示意，RTX 2000 一张显卡的价格接近左图PC的一半。

不能加DRAM，只能从模型大小入手，先得分析模型实际工作流。

LLM 推理工作流

LLM 从 input 到 output 工作流程

模型架构：

输入经过嵌入层后，进入一个循环18次的模块。该模块包含两个主要部分：

MLP层（Gelu激活函数）和RMS归一化。
注意力层和RMS归一化。

最后，经过线性层和softmax层，输出模型的概率结果。

嵌入层/注意力层/MLP层模型体积占比

嵌入层（Embedding Layer）占模型尺寸的 20%。
注意力层（Attention Layer）占 8%。
MLP层占 72%，并且该层是稀疏的。

Note：从工作原理来看，压缩模型体积的重点工作应关注嵌入层和MLP层。

常用数据驻留DRAM策略

每个LLM的层组件组成百分比各不相同。
基于层的组成结构，我们可以让部分LLM驻留在GPU的VRAM中。

例子：在Gemma 2B参数模型中，保持28%的部分（其中20%为嵌入层（Embedding Layer），8%为注意力层（Attention Layer））驻留在GPU上，而72%的MLP层按需加载（非驻留部分）。

图展示了一种优化的方式，将LLM的部分层（如嵌入层和注意力层）驻留在GPU的VRAM中，而其余的MLP层则根据需求从SSD加载。这种方法在避免占用大量GPU显存的同时，依然可以实现高效的推理计算。这种混合存储架构利用了GPU的高效性与SSD的容量优势。

从Flash中加载运行模型参数

从NVMe设备加载大型语言模型

通过NVMe存储设备加载大型语言模型（LLM）的流程，以及不同步骤的耗时。

GPU单元：

GPU负责处理嵌入层和注意力层，处理时间为 0.2毫秒/层。
GPU通过CUDA内核API与CPU单元进行通信，每层的处理时间为 18毫秒/层。

CPU单元：

运行LLM的应用程序通过XNVMe（增强型NVMe接口）访问数据。
CPU通过PCIe与SSD进行通信。

SSD单元：

SSD上存储了Gemma-2B完整模型。
从SSD加载每层的时间为 63毫秒/层，使用XNVMe接口后时间减少到 24毫秒/层，实现了大约 60%的减少。

展示了从NVMe设备加载LLM的具体流程，并强调了XNVMe技术对加载效率的提升。在未使用XNVMe时，SSD的加载时间为63毫秒/层，而使用XNVMe后时间降低至24毫秒/层，减少了60%。GPU处理速度较快，但整体加载性能受限于数据从SSD传输到CPU再到GPU的速度。因此，XNVMe技术的应用显著提升了模型的加载效率。

Note: XNVMe 技术关键在于优化数据在操作系统的IO过程，经典路径是需要用户层-内核态调度，而 XNVME 的核心贡献是用户态直接调用，从而达到加速的目的。

WD作为传统存储器厂商，站在Flash特性上研究推理优化，别忘了还有个CXL，其在通信技术上的布局不仅仅是IO路径优化，而是改变现有服务器通信基座。

更进一步，从Flash到GPU？

换一个视角：数据能不能从Flash直接到GPU的显存？ 跳过CPU的传输过程。

通过从闪存中流式加载参数到VRAM，同时仍然保持可接受的推理性能，并提出了一些优化方法。

问题：

我们能否通过从闪存中流式加载参数到VRAM，同时保持可接受的推理性能？
许多LLM（大型语言模型）具有高度稀疏性。我们能否利用这一点，有选择地加载参数，以避免冗余计算？

优化策略：

Squeeze（压缩）：

LRP（Low Rank Predictor）：预测哪些神经元将保持活跃，哪些将被置零；然后省略掉那些被置零的神经元。

Speed Up（加速）：

Row Column Bundling（行列绑定）：将上下投影神经元进行聚类。这有助于减少从SSD读取的次数。

为什么LLM的稀疏性被加以利用后，能避免冗余计算？

LLM（大型语言模型）的稀疏性指的是模型中的一部分神经元或参数在特定的计算过程中并不总是被激活或使用。例如，某些输入在通过模型时，只有一部分神经元会参与计算，而另一部分神经元则处于“非活跃”状态，这意味着这些非活跃神经元在当前推理任务中并没有贡献输出。这种现象被称为稀疏性。

利用稀疏性避免冗余计算的关键在于以下几点：

1. 减少不必要的计算：

   稀疏性意味着不是每个神经元在每一步推理中都需要计算输出。如果我们可以预测哪些神经元不会参与当前的推理过程（例如，通过低秩预测器，LRP），那么就可以跳过对这些神经元的计算，直接忽略它们。这样可以减少不必要的计算，节省计算资源和时间。

2. 选择性加载参数：

   如果某些神经元的权重在推理过程中不被激活，就可以不加载这些权重，特别是在存储资源有限的设备上（如GPU的VRAM）。通过有选择性地加载参数，仅加载那些在当前推理任务中需要的参数，可以有效减少从存储设备（如SSD）到GPU的读取操作。这种方法不仅减少了数据传输的开销，也减少了对内存空间的占用。

3. 加速推理过程：

   跳过冗余计算和减少不必要的参数加载直接提高了推理速度。模型中活跃的神经元只占一部分，稀疏性使得我们能够专注于少数真正需要参与运算的部分，大大减少了总的计算负荷，从而加速模型的推理。

4. 降低硬件压力：

   在实际应用中，GPU和CPU的资源有限，特别是在客户端设备上。稀疏性允许模型避免不必要的内存使用和计算，优化硬件资源的使用。这在LLM模型规模巨大、设备硬件资源有限的情况下显得尤为重要。

Note：模型稀疏性研究是推动其在有限资源端、边设备运行的关键！

存储硬件或软件厂商，能在模型稀疏性上尝试哪些创新？

硬件厂商

专用加速器： 开发专门针对稀疏矩阵运算优化的硬件加速器。这些加速器可以设计有特殊的电路来跳过零值计算，从而节省处理时间和能耗。

内存架构： 采用能够更高效地存储稀疏数据结构的内存架构，例如压缩感知内存（Compressed Sensing Memory, CSM）或稀疏RAM，它们可以在物理层面上减少对零值的存储需求。

软件厂商

稀疏数据格式： 提供高效的数据格式来表示稀疏张量，如CSR（压缩行存储）、CSC（压缩列存储）或COO（坐标列表），并为这些格式提供高性能的库函数。

优化算法： 开发专为稀疏数据优化的线性代数库，包括稀疏矩阵乘法、解方程组等核心操作。