苹果公司在大型语言模型(LLM)方面入局较晚,但最近正在采取一些有趣的举措。该公司的研究人员最近推出了一种方法,可以显著降低运行LLM所需的内存。这一突破依赖于新颖的存储和内存管理技术,这些技术可以在保持令人印象深刻的低延迟的同时,动态地在闪存和DRAM之间传输模型的权重。
他们的方法使得即使部分加载到设备RAM上,这些大型模型也可以以相当的吞吐量运行。这对于苹果这样没有超级计算服务业务的公司来说可能是重要的,因为它将从设备上的AI中获得巨大的好处。这项技术可以使未来的iPhone和Mac利用LLM的能力,而不会压垮系统的内存。
LLM的内存要求
LLM因其巨大的内存需求而声名狼藉。例如,一个70亿参数的模型,以半精度要求超过14GB的RAM。这样的需求超出了大多数边缘设备的容量。
解决这一问题的一种方法是量化,通过将参数转换为8位或4位整数来压缩模型。然而,量化通常需要对模型进行修改,甚至需要进行全面重新训练。在某些情况下,量化后的LLM仍然太大,无法在设备上进行推断。
苹果的研究试图从一个不同的角度解决问题:我们如何在缺乏必要内存的硬件上部署模型?
从闪存加载LLM
解决内存瓶颈的一种策略是将LLM存储在闪存中,并为推断任务逐渐加载到RAM中。虽然设备上的闪存比DRAM更丰富,但它至少慢了一个数量级。
使用闪存进行天真的推断可能需要对每次前向传递重新加载整个模型。这个过程很慢——即使是对压缩模型,通常也需要几秒钟——而且耗能大。
在他们的论文中,苹果的研究人员介绍了一系列旨在简化从闪存到DRAM加载模型的过程的优化技术。这些方法旨在在保持内存消耗低的同时维持推断的速度。
利用模型稀疏性
语言模型依赖于变换器块,它由两个关键元素组成:注意力机制和前馈(FFN)网络。研究表明,LLM中的FFN具有高度稀疏性——许多值在激活后变为零或接近零,对推断来说无关紧要。令人惊讶的是,一些FFN层的稀疏度可以超过95%。
苹果的研究人员建议利用这种稀疏性来优化模型推断。通过发现并仅在推断期间加载非稀疏元素,内存成本可以大幅降低。他们的策略涉及完全加载比较小的注意力层,同时选择性地仅将FFN的非稀疏部分加载到DRAM中。
“我们的方法包括在推断过程中迭代地仅将必要的非稀疏数据从闪存转移到DRAM进行处理。”研究人员写道。
为了实现这一点,他们采用了一个“低秩预测器”网络,用于确定FFN的哪些部分将是非稀疏的,并减少需要加载的数据量。然后将预测器识别的活跃神经元转移到内存中。
滑动窗口
苹果的研究人员还设计了一种“滑动窗口技术”,用于管理在模型推断期间的神经元加载与卸载。这种方法仅将最近一小部分输入标识的神经元数据保持在内存中,当新标识进来时,释放旧标识的空间。他们的研究揭示,随着每个新标识的出现,只需要交换极少数神经元,从而简化了从闪存到RAM的数据传输。
“这种策略允许有效地利用内存,因为它释放了之前分配给不再处于滑动窗口中的旧标识的神经元数据的内存资源。”研究人员写道。
滑动窗口的大小可以根据运行模型的设备的内存限制进行配置。
优化存储和加载
苹果团队的研究表明,不同层之间的神经元激活遵循可以用于存储优化的模式。他们建议根据需要一起加载的神经元重新组织模型权重的存储。这种安排允许系统在一次读取操作中提取一个连续的数据块。
研究人员还研究了一起存储共激活神经元的好处。这种方法将允许系统从一个大的闪存段提取数据。尽管来自这个特定实验的结果未达到预期,但这些洞察力为未来的研究方向提供了指导。
此外,团队还设计了一种新颖的数据结构和动态内存分配算法。这项创新减少了内存分配的频率,并允许模型无显著性能下降地访问约25%的FFN神经元。
在低内存设备上加速LLM
苹果的研究团队在两个不同的模型上测试了他们的技术,分别是OPT 6.7B和一个Falcon 7B的稀疏化变体,并在两个硬件设置上进行了试验。第一个平台是配备1TB SSD作为闪存的Apple M1 Max。第二组测试是在配备24GB NVIDIA GeForce RTX 4090显卡的Linux机器上进行的。在他们的实验中,他们在任何给定时间只加载了模型权重的约50%。
在M1 Max上,每次推断从闪存到RAM的加载模型将导致每个标识的延迟为2.1秒。然而,通过实施他们的新颖技术(稀疏性预测、窗口化和智能存储)后,该团队将这个延迟减少到大约200毫秒。在配备GPU的系统上的改进更加明显。
“我们已经证明了能够运行LLM,这些模型的大小是可用DRAM的两倍,与传统的加载方法相比,我们在CPU上实现了4-5倍的推断速度加速,在GPU上实现了20-25倍。”研究人员写道。
研究人员认为,他们的发现可能对未来的研究有重要影响,并强调了“在开发推断优化算法中考虑硬件特性”的重要性。
“我们相信,随着LLM在大小和复杂性上的持续增长,像这项工作这样的方法将是在广泛的设备和应用程序中充分利用它们潜力的关键。”他们写道。