Pie: Pooling CPU Memory for LLM Inference

大语言模型（LLM）的快速发展已经彻底改变了自然语言处理和人工智能分析，但它们日益增长的规模和内存需求带来了重大挑战。一种常见的解决方案是溢出到CPU内存；然而，传统的GPU-CPU内存交换通常会导致更高的延迟和更低的吞吐量。 本文介绍了一种名为Pie的大语言模型推理框架，该框架通过性能透明的交换和自适应扩展来解决这些挑战。通过利用可预测的内存访问模式和现代硬件（如NVIDIA GH200 Grace Hopper超级芯片）的高带宽，Pie能够在不影响前台计算的情况下实现并发数据交换，从而在不增加延迟的情况下扩展有效内存。自适应扩展根据实时信息动态调整CPU内存分配，优化不同条件下的内存使用和性能。 Pie保持了低计算延迟、高吞吐量和高弹性。我们的实验评估表明，Pie在缓存预热期间实现了最优的交换策略，并有效地平衡了增加的内存容量与对计算的微小影响。凭借其扩展容量，Pie在吞吐量上比vLLM高出1.9倍，在延迟上高出2倍。此外，Pie可以在保持相同性能的同时将GPU内存使用量减少多达1.67倍。与基于离线配置文件的交换解决方案FlexGen相比，Pie实现了显著降低的延迟和9.4倍的更高吞吐量。

该论文旨在解决大规模语言模型（LLMs）在进行自然语言处理和AI分析时面临的内存需求高、传统GPU-CPU内存交换导致的高延迟和低吞吐量的问题。这是一个随着LLMs规模不断增大而日益突出的问题。

论文提出了一种名为Pie的LLM推理框架，通过性能透明的内存交换和自适应扩展技术，利用可预测的内存访问模式和现代硬件（如NVIDIA GH200 Grace Hopper Superchip）的高带宽，实现了并发数据交换而不影响前台计算，从而扩展有效内存容量且不增加延迟。这一方法在现有研究基础上，提供了更高效、更灵活的内存管理方案。

实验设计方面，论文通过与vLLM和FlexGen等现有解决方案的对比，展示了Pie在吞吐量、延迟和GPU内存使用方面的显著优势。具体来说，Pie在吞吐量上比vLLM高出1.9倍，在延迟上减少2倍，并能将GPU内存使用降低1.67倍。此外，Pie还支持动态调整CPU内存分配，以优化不同条件下的性能。论文提供了详细的实验设置和结果分析，验证了Pie的有效性。目前，Pie的源代码已开源，为后续研究提供了基础。

近年来，针对LLM的内存管理和优化，已有多个相关研究。例如，《Optimizing Large Language Model Inference with Memory-Constrained Devices》探讨了在资源受限设备上优化LLM推理的方法；《Memory-Efficient Fine-Tuning of Large Language Models》则研究了如何在有限内存下对LLM进行微调。此外，《FlexGen: An Efficient and Flexible Framework for Large-Scale Language Model Inference》提出了基于离线配置文件的内存交换策略，但与Pie相比，其在延迟和吞吐量方面表现较差。