在构建基于RAG（检索增强生成）的问答系统时，我们经历了一次典型的内存瓶颈问题，并通过Python生成器巧妙地化解了危机。本文将从整体架构、问题复现、解决方案及最终效果四个方面，完整记录这一过程，希望能为遇到类似问题的开发者提供参考。

一、系统架构概览

我们的RAG服务基于FastAPI + LangChain 1.0+ 构建，核心流程如下：

1. 文档索引阶段

• 递归切片：对原始文档按标题、段落和语义边界进行递归分割，保证语义完整性

• 双路向量化：使用同一Embedding模型生成1024维稠密向量和稀疏向量（基于词袋或BM25特征）

• 存储：将稠密向量、稀疏向量和原始文本块一并存入Milvus 2.x

2. 检索增强阶段

• 用户查询：对输入问题同样生成稠密+稀疏向量

• 混合检索：

  • 稠密向量检索（ANN）：返回20条候选

  • BM25稀疏检索：返回20条候选

• RRF融合：根据权重分配（稠密0.7，BM25 0.3）重新排序，取前20条

• 精排重排：使用Cross-Encoder风格的Rerank模型（如BGE-reranker）对20条结果去重并打分，筛选出最相关的10条

• LLM生成：取最终分数最高的3条作为上下文，调用大模型生成答案

整个系统在设计上追求高召回与高精度的平衡，特别是在混合检索和Rerank环节投入了较多优化。

二、问题出现：切片向量化时的内存“洪峰”

在索引阶段，我们最初采用的实现方式是将所有切片一次性读入内存，批量调用Embedding模型。文档量不大时（几千个切片以内）运行平稳。但当测试语料上升到数十万份长文档时，递归切片产生了超过200万个文本块。

每个文本块平均200 tokens，再加上对应的稠密向量（1024维float32 ≈ 4KB）和稀疏向量（存储为字典或list，通常10~20KB），单纯这批切片对象的内存占用就轻松突破了20GB。更严重的是，LangChain的batch调用会在内部积累所有输入和输出，导致内存峰值达到原始数据量的3~5倍。最终结果就是OOM Killer无情地终止了进程。

初步排查发现，根本原因在于数据流水线上缺乏惰性求值机制——切片生成、向量化、写入Milvus全部是”全量加载，批量处理“的模式。

三、灵机一动：生成器重构内存流水线

解决思路并不复杂：将“全量加载”改为“流式处理”。Python的生成器（yield）天然适合这种场景。具体改造如下：

1. 递归切片器变成生成器

python

def recursive_split_generator(documents, chunk_size, overlap):
    for doc in documents:
        # 递归分割逻辑（略）
        for chunk in split_recursively(doc, chunk_size, overlap):
            yield chunk  # 每生成一个切片就吐出，不等待

2. 逐批向量化与写入

不再一次性收集所有切片，而是维护一个小批量缓冲区（例如batch_size=64）：

python

def index_pipeline(doc_source):
    buffer = []
    for chunk in recursive_split_generator(doc_source):
        buffer.append(chunk)
        if len(buffer) >= BATCH_SIZE:
            # 编码当前批次
            dense_embs = embedding_model.encode(buffer)  # shape: [64, 1024]
            sparse_embs = sparse_encoder.encode(buffer)
            # 写入Milvus
            milvus_client.insert(collection, buffer, dense_embs, sparse_embs)
            buffer.clear()
    # 处理最后遗留的不足一批的数据
    if buffer:
        ...

3. 关键点：生成器带来的双重红利

• 内存占用扁平化：任何时候内存中只保留batch_size（如64）个切片及其向量，不再存储200万个对象。内存从20GB+骤降至约500MB（含模型常驻开销）。

• CPU/GPU利用率提升：以前是“干等内存分配，然后批量计算”，现在是不断生产-消费，GPU编码和Milvus写入可以流水线重叠（若使用异步队列则更佳）。

四、效果与注意事项

实际效果

• 内存峰值：从22GB降至1.2GB（含模型权重约800MB）

• 吞吐量：原方案处理100万切片约需120分钟（大部分时间在GC和swap），新方案仅需40分钟

• 系统稳定性：长时间运行无OOM，Milvus连接池正常复用

踩坑提醒

1. Milvus批量插入限制：单次插入不宜超过512条，建议配合生成器的batch控制。

2. 稀疏向量的序列化：Milvus 2.2+支持sparse_vector类型，直接传入dict[int, float]即可，无需手动压缩。

3. 生成器中的递归深度：递归切片本身可能深度较大（如嵌套标题），注意Python默认递归限制，必要时改用栈实现。

4. LangChain 1.0+ 的兼容性：新版本的BaseDocumentTransformer支持流式接口，可以直接封装生成器。

五、总结与展望

通过一个简单的生成器模式，我们不仅解决了RAG索引阶段的内存爆炸问题，还顺便提升了整体吞吐量。这再次印证了在处理大规模数据时惰性求值的重要性。

当然，生成器只是众多优化手段之一。随着数据量进一步增长（千万级切片），我们可能需要引入：

• 分布式任务队列（Celery + Redis）

• 更加激进的稀疏向量压缩（如量化）

• Milvus的批量导入工具（bulk_insert）

但至少现在，我们的RAG服务可以在单台32GB内存的机器上平稳运行，为上层应用提供了可靠的检索基础。如果你也在构建类似系统，不妨从审视数据流水线开始——往往最简单的yield就能解决最头疼的问题。

---

作者：琦琦
技术栈：FastAPI, LangChain 1.0+, Milvus, Python, BGE/OpenAI Embedding, Reranker