基于SSD的RAG技术方案，推动LLM规模扩展

news2024/9/23 5:56:11

随着大型语言模型（LLM）的不断发展，它们在虚拟助手、聊天机器人和对话系统等应用中发挥着重要作用。然而，LLM面临的挑战之一是它们可能会生成虚假或误导性的信息，即所谓的“幻觉”。为了解决这一问题，检索增强生成（RAG，Retrieval Augmented Generation）作为一种解决方案应运而生，它允许LLM从外部可信的数据源中提取信息，以验证和支撑其输出内容的真实性。

RAG的基本原理是让LLMs在生成回答时，能够从外部数据集中搜索与查询相关的文档，从而获取更准确和可靠的信息。具体来说，当用户提出问题时，LLMs会同时利用自身的大规模训练数据集和外部数据集进行查询。外部数据集的向量嵌入（如私有、可信、最新的文档）会被存储起来，通过近似最近邻搜索（ANNS，Approximate Nearest Neighbor Search ）算法，快速检索与查询最相关的信息元素，为LLMs的回答提供支持。

有效的ANNS算法是RAG能否成功的关键。为了快速准确地检索到与查询相关的信息，需要使用ANNS算法。

目前，Hierarchical Navigable Small World（HNSW）是领先的内存型ANNS算法，但它存在着无法经济地扩展的问题。

随着RAG数据集规模的增大，尤其是向量嵌入的高维度（768D - 1536D），HNSW需要大量的DRAM来存储向量和索引，这导致了高昂的成本，限制了RAG数据集的大小和其基础有效性。

与此相反，SSD-based ANNS算法，如微软开发的DiskANN，则是基于SSD的解决方案。DiskANN通过将索引存储在SSD上，而将量化后的向量（Product Quantization, PQ）保存在内存中，针对SSD实现算法优化和优化SSD的访问模式，从而实现了在成本效益方面的显著提升，同时保持了与内存基方案相媲美的性能。

铠侠在FMS2024上展示了一种适用于大规模、高维度RAG数据集的DiskANN。他们使用铠侠CD8P系列PCIe NVMe数据中心SSD对DiskANN进行基准测试。结果显示，基于SSD的DiskANN与内存中的HNSW提供的性能相当。

DiskANN显著降低了HNSW的DRAM占用，减少了高达89%的DRAM footprint。这使得在提供可比性能的同时，能够大幅降低成本，为RAG应用在大规模场景下的实现提供了可能。

向量数据库从对象存储中加载向量和索引数据到查询节点。这意味着数据的存储和管理是分离的，对象存储负责数据的长期保存，而查询节点在需要时从对象存储中获取数据进行处理。在大规模RAG应用中，查询节点使用SSD作为搜索介质。这是因为SSD具有较高的读写速度和性能，可以更快地处理大规模的数据搜索请求，从而提高RAG应用的效率和响应速度。

使用DiskANN的相同架构能够实现与HNSW相当的吞吐量。这意味着DiskANN在数据搜索和处理能力方面能够达到与传统的HNSW算法相近的水平，保证了RAG应用的性能。更为重要的是，DiskANN显著降低了查询节点的成本。这是通过优化算法和存储架构实现的，使得在不牺牲性能的前提下，能够减少硬件资源的消耗和成本支出。

SSD在RAG-Based LLMs的向量数据库架构中发挥着关键作用，它为大规模RAG应用提供了高效的搜索介质，而DiskANN的应用则在保证性能的同时，大幅降低了查询节点的成本，推动了RAG应用在大规模场景下的可扩展性和经济性。这种架构的优势在于能够更好地满足RAG应用对数据处理速度和成本控制的需求，为LLMs的发展和应用提供了有力的支持。

为了进一步推动RAG技术的应用和发展，KIOXIA提出了RAG优化的SSD解决方案（RAG Optimized SSD Solution, ROSS）。ROSS的关键点在于将向量数据库从DRAM迁移到SSD，这样不仅可以降低成本，还能提高系统的可扩展性。KIOXIA的ROSS概念利用SSD来承载ANNS，其性能与DRAM解决方案相当，但不受DRAM大小限制，因此可以支持更大规模的向量数据库。

KIOXIA AiSAQTM是KIOXIA为ROSS开发的首个成果，它使向量数据库能够从对象存储中加载向量和索引数据到查询节点，并使用SSD作为大规模RAG应用的搜索介质。

KIOXIA AiSAQTM（All-in-Storage ANNS with Product Quantization）能够实现从“大部分在存储中”（如DiskANN）到全在存储中、几乎零DRAM的架构转变。KIOXIA AiSAQTM将进一步扩展RAG的规模，使RAG应用能够达到几乎无限的规模。

综上所述，RAG和ANNS是现代LLM解决方案的重要组成部分。然而，基于内存的ANNS算法由于高昂的成本和有限的扩展能力，无法满足日益增长的需求。借助SSD技术的进步，特别是像KIOXIA提出的ROSS和AiSAQ这样的创新解决方案，我们可以预见未来RAG技术将能够更广泛地应用于实际场景中，实现更大规模的数据处理和服务提供。随着这些技术的发展，我们可以期待更加智能、高效的语言模型应用，为用户提供更加精准、可靠的服务体验。

参考文献：FMS 2024-Kioxia-《Flash is Driving Scale in RAG-Based LLMs》

如果您看完有所受益，欢迎点击文章底部左下角“关注”并点击“分享”、“在看”，非常感谢！

精彩推荐: