FlashRAG-Paddle:基于 PaddleNLP 的高效开发与评测 RAG 框架
论文标题:
FlashRAG: A Modular Toolkit for Efficient Retrieval-Augmented Generation Research
作者机构:
中国人民大学高瓴人工智能学院
论文链接:
https://arxiv.org/abs/2405.13576
项目链接 (Paddle 版本):
https://github.com/RUC-NLPIR/FlashRAG-Paddle
数据集链接:
https://huggingface.co/datasets/RUC-NLPIR/FlashRAG_datasets
产品亮点
1. FlashRAG-Paddle:组件化、模块化的 RAG 框架
全面且可定制的 RAG 框架,集成检索器、重排序器、生成器和压缩器等核心组件,提供 36 个基准数据集和 9 种先进算法,辅以高效预处理脚本,简化流程,轻松实现复杂 RAG 场景下的模型测试与验证。
2. PaddleNLP:超大 Batch 嵌入表示学习和多硬件高性能推理
PaddleNLP 提供一站式大语言模型解决方案,支持超大 Batch 嵌入学习,多硬件高性能推理,涵盖了 INT8/INT4 量化技术,以及 PageAttention、FlashDecoding 等高效的注意力机制优化和 TensorCore 深度优化,从而大幅提升训练与推理效率,全方位满足多样化的应用需求。
3. FlashRAG & PaddleNLP:检索增强生成结合高性能推理,更准更快,提升用户体验
基于飞桨框架 3.0 版本,PaddleNLP 内置了全环节算子融合等技术,使得 FlashRAG 推理性能相较于 transformers 动态图推理实现了 70% 以上的显著提升,结合检索增强知识,输出结果更加准确,为 FlashRAG 框架使用者带来了敏捷高效的使用体验。

背景介绍
在大语言模型时代,检索增强生成已成为缓解幻觉问题的有效解决方案。通过外接知识库,大模型可以在长尾问题、垂域问题上有良好的表现。检索增强生成(RAG)技术的广泛应用和巨大潜力吸引了大量研究关注。
然而,近年来随着大量 RAG 算法和模型的出现,如何在一致的环境下快速构建并比较评估这些方法变得越来越具有挑战性。同时,如何利用大语言模型为文本更快更好的构建准确嵌入表示、如何加速推理生成速度等方面也日益受到关注。
为解决上述问题,中国人民大学高瓴人工智能学院联合百度共同发布了 FlashRAG-Paddle 框架,其中内置 36 个经过预处理的 RAG 数据集以及 9 种预实现的 RAG 算法,能够帮助研究人员高效地在 RAG 领域进行复现、基准测试和开发新算法。
同时,借助基于飞桨框架 3.0 版本打造的 PaddleNLP 大语言模型套件,通过极致的全流程优化,为 RAG 中的检索器、生成器、重排器、精炼器提供从组网开发、预训练、精调对齐、模型压缩以及推理部署的一站式解决方案。

FlashRAG-Paddle 框架
FlashRAG-Paddle 是一个用于检索增强生成(RAG)研究的 Python 工具包,它基于国产深度学习平台飞桨框架和 PaddleNLP 大语言模型套件构建,并针对国产芯片进行了优化。
- 全面且可定制的框架: 包含 RAG 场景所需的核心组件,如检索器、重排序器、生成器和压缩器,可灵活组合成复杂的流程。
- 综合基准数据集: 包含 36 个经过预处理的 RAG 基准数据集,用于测试和验证 RAG 模型的性能。
- 预实现的先进 RAG 算法: 提供 9 种先进的 RAG 算法,可在不同设置下轻松复现结果。
- 高效的预处理阶段: 通过提供各种脚本,如用于检索的语料库处理、检索索引构建以及文档的预检索,简化了 RAG 工作流程的准备工作。
为了能够灵活高效地搭建各类 RAG 系统,FlashRAG 采用了组件化、模块化的设计理念,将整个框架的设计包括三个层面:组件层、流程层、数据层。

2.1 组件层
组件层位于框架的最底层,提供了构建 RAG 系统所需的各类基础组件,主要包括:
- 检索器 (Retriever): 负责从知识库中检索与查询最相关的文档。FlashRAG 支持使用 embedding 模型以及基于词项匹配的 BM25 方法等。
- 生成器 (Generator): 根据给定的文本(通常是查询和检索结果的拼接)生成最终的回复。框架支持使用各类 LLM 模型,并支持 FastChat、vllm 等加速方案。
- 重排器 (Reranker): 对检索结果进行重新排序,以进一步提升与查询的相关性。重排器分为交叉编码器和双塔编码器两种类型。
- 精炼器 (Refiner): 对输入的文本进行进一步的精炼和压缩,去除冗余信息。目前已支持抽取式、生成式、基于 LLMLingua 和基于 Selective-Context 的多种精炼器。
基于实现的组件库,用户可以自由选择需要使用的组件来完成自己的特定需求。
2.2 流程层
流程层位于组件层之上,通过组装各类组件实现端到端的 RAG 流程。基于各种方法的推理路径,我们将 RAG 流程分为了四大类:
- Sequential: 顺序执行 retriever、refiner、reranker、generator 等组件,是最基础的 RAG 流程。
- Conditional: 通过 judger 模块判断不同类型的查询,并选择不同的执行路径
- Branching: 并行执行多条路径,并将各路径的生成结果进行整合,代表工作如 REPLUG、SuRe 等。
- Loop: 通过迭代的方式交替执行 retriever 和 generator,代表工作如 Self-Ask、Self-RAG、FLARE、IRCoT 等。

2.3 数据层
最上层为数据层,包括用于检索的语料数据以及用于评估的各种任务数据。在 RAG 流程运行完成后,会自动计算相关的评价指标并保存评测结果。
我们收集并处理了 RAG 研究中广泛使用的 35 个数据集,并对其进行了预处理,以确保格式一致,便于使用。对于某些数据集,我们根据社区中常用的方法对其进行了调整以满足 RAG 任务的要求。所有数据集均可在 Huggingface 平台上进行下载。

PaddleNLP:超大 Batch 嵌入表示学习和多硬件高性能推理助力检索增强生成
PaddleNLP 是基于飞桨框架打造的大语言模型套件,通过极致的全流程优化,为开发者提供从组网开发、预训练、精调对齐、模型压缩以及推理部署的一站式解决方案。
- 模型组网简化与参数多样: PaddleNLP 通过统一分布式表示与自动并行技术,显著简化了组网开发的流程,减少了分布式核心代码量 50% 以上,结合多种并行策略 Llama 3.1 405B 等超大规模模型能够开箱即用。此外,PaddleNLP 预置了 80 多个主流模型的训练、压缩、推理全流程方案,满足了不同应用场景下的多样化需求。
- 精调与对齐性能提升: 借助飞桨框架独有的 FlashMask 高性能变长注意力掩码计算机制和 Zero Padding 零填充数据流优化技术,PaddleNLP 有效减少了无效数据填充带来的计算资源浪费,显著提升了精调和对齐的性能。以 Llama 3.1 8B 模型为例,其性能相较于 LLaMA-Factory 方案实现了 1.2 倍的提升,单机即可轻松完成 128K 长文的 SFT/DPO 任务。
- 硬件适配广泛与高效: PaddleNLP 基于飞桨插件式松耦合统一硬件适配方案(CustomDevice),仅需适配 30 余个接口即可实现大模型的基础适配,支持英伟达 GPU、昆仑芯 XPU、昇腾 NPU、燧原 GCU 和海光 DCU 等多款主流芯片的大模型训练和推理。依托框架提供的多种算子接入模式和自动并行调优技术,PaddleNLP 实现了框架与芯片间的软硬协同性能优化,为用户提供了更加高效、稳定的模型训练和推理体验。
3.1 超大 Batch 嵌入表示学习
嵌入表示学习采用 In-batch negative 策略,即将 batch 内的所有其他样本视作负样本,定义 Contrastive loss 函数如下:

其中是第 i 个文本和第 j 个文本之间的余弦相似度,s_ii 是匹配样本的相似度。从公式中可以看出随着 batch size 的增大,模型能接触并学习更多负样本,进而提取出更具判别性的特征。因此在大模型训练过程中,需要更大的 Batch 提升模型性能。
PaddleNLP 通过在数据并行中支持 In-batch negative 策略以及 flashmask 等显存优化策略,在训练 embedding 模型时,提升每次训练的 batch size,进而显著提高训练效率和效果。
为了扩大全局 batch,PaddleNLP 采用 gradient cache 策略。该策略首先执行模型前向传播,计算 micro-batch 输出并舍弃中间隐藏状态。随后,累计这些输出以计算 Contrastive loss,实现超大 batch 的对比学习效果。最后,针对各 micro-batch 重新计算隐藏状态,进行前向与反向传播,并依据对比学习梯度优化模型参数来提升训练效果。
3.2 多硬件高性能推理
特点一:高性能推理优化,支持 Llama 3.1 405B 高性能推理
PaddleNLP 构建了大模型高性能推理方案,支持 Llama 3.0、Llama 3.1、Mixtral 等一系列大语言模型推理。当前支持技术包含:
- Weight Only INT8 及 INT4 推理,支持权重、激活、Cache KV 进行 INT8、FP8 量化的推理;
- 注意力机制支持 PageAttention、FlashDecoding 等优化;
- 支持基于 TensorCore 深度优化。

PaddleNLP 高性能推理通过内置全环节算子融合策略,获取更优推理性能。在 7B、14B、32B 和 72B 模型的推理性能上,PaddleNLP 后端相比 transformers 后端动态图推理提速 70% 至 119%。同时,Llama 3.1 405B 作为开源社区中的最大模型,PaddleNLP 推理也快速支持了此模型,单机 8 卡即可实现快速推理能力。
特点二:多硬件大模型推理支持

在硬件支持方面,飞桨推理引擎展现出了强大的兼容性和灵活性。当前在支持英伟达 GPU 的基础上,还支持了国产及国际领先的 AI 芯片生态,包括昆仑 XPU、昇腾 NPU、海光 DCU、燧原 GCU、英特尔 X86 CPU 等多种硬件的大模型推理,不同硬件的推理入口保持统一,仅需修改 device 即可支持不同硬件推理。
PaddleNLP 的这一特性强化了 RAG 在更广泛的多硬件环境下的部署场景,满足不同用户的多样化需求,推动 RAG 在更多实际业务中的应用与落地。

FlashRAG + PaddleNLP 快速打造 RAG 文档问答应用
RAG(检索增强生成)技术巧妙融合了信息检索与先进的生成模型,通过从丰富的外部知识库中精准检索相关信息,为大型语言模型提供强有力的回答辅助。这一创新技术充分发挥了信息检索的精准性与生成模型的创造力,能够自动生成既高质量又准确,同时紧密贴合上下文的文档总结。接下来,我们将为您详细介绍如何仅需三步,即可轻松构建 RAG 文档总结应用。

4.1 构建自己的语料数据库
步骤 1:准备语料
首先需要准备预料库来作为检索召回的基础,预料库保存为 jsonl 以下格式:
{"id": "0", "contents": "contents for building index"}
{"id": "1", "contents": "contents for building index"}
更多语料制作细节,请前往 https://github.com/RUC-NLPIR/FlashRAG-Paddle 查看。
步骤 2:索引
然后,使用以下代码构建您自己的索引。
对于密集检索方法,通过使用 embedding 模型,将文档嵌入到密集向量中,之后使用 faiss 集向量索引。
对于稀疏检索方法,基于 Pyserini 或 bm25s,依据词频将语料构建为 Lucene 的倒排索引。
密集检索方法
python -m flashrag.retriever.index_builder \
retrieval_method e5 \
model_path intfloat/e5-base-v2/ \
corpus_path indexes/sample_corpus.jsonl \
save_dir indexes/ \
use_fp16 \
max_length 512 \
batch_size 256 \
pooling_method mean \
faiss_type Flat
稀疏检索方法(BM25)
python -m flashrag.retriever.index_builder \
--retrieval_method bm25 \
--corpus_path indexes/sample_corpus.jsonl \
--bm25_backend bm25s \
--save_dir indexes/
4.2 配置检索器和生成器
步骤 1:配置检索模型和生成模型
通过这段代码,我们指定了检索模型和生成模型的具体设置,为 RAG 运行提供了必要的配置信息。
config_dict = {
"save_note": "demo",
"model2path": {
"e5": "intfloat/e5-base-v2",
"llama3-8B-instruct": "meta-llama/Meta-Llama-3-8B-Instruct"
},
"retrieval_method": "e5",
"generator_model": "llama3-8B-instruct",
"corpus_path": "indexes/general_knowledge.jsonl",
"index_path": "indexes/e5_Flat.index",
}
config = Config("my_config.yaml", config_dict=config_dict)
步骤 2:定义对话模板
为了让 AI 助手更加人性化,我们需要为它设定一些对话模板。这些模板会指导 AI 如何回应用户的输入,并确保回复既友好又准确。
system_prompt_rag = (
"你是一个友好的 AI 助手。"
"像人类一样回应输入,如果输入中有指令,请遵循指令进行回应。"
"\n以下是一些提供的参考信息。你可以使用这些信息来回答问题。\n\n{reference}"
)
base_user_prompt = "{question}"
prompt_template_rag = PromptTemplate(config, system_prompt=system_prompt_rag, user_prompt=base_user_prompt)
步骤 3:加载检索模型和生成模型
在配置好对话模板之后,接下来我们需要加载之前配置的检索模型和生成模型。这两种模型在 PaddleNLP 库中已经集成,所以我们可以直接调用它们,无需从头开始训练或编写复杂的代码。
retriever = load_retriever(config)
generator = load_generator(config)
4.3 检索增强,生成输出
通过这段代码,RAG 首先使用检索模型根据用户的查询找到最相关的文档。然后,它将这些文档信息整合到对话模板中,形成一个包含引用信息的输入提示。最后,使用生成模型根据这个输入提示来生成智能化的回应,并打印出来。
retrieved_docs = retriever.search(query, num=topk)
input_prompt_with_rag = prompt_template_rag.get_string(question=query, retrieval_result=retrieved_docs)
response_with_rag = generator.generate(
input_prompt_with_rag, temperature=temperature, max_new_tokens=max_new_tokens
)[0]
print(response_with_rag)

总结
FlashRAG-Paddle 这一工具包不仅有助于现有 RAG 技术的复现,还支持新方法的开发,更重要的是,它为基于国产硬件和软件的 AI 应用提供了更广泛、更灵活的开发和部署选项,从而推动了 AI 技术的自主创新和国产化进程。
FlashRAG 和 FlashRAG-Paddle 的所有资源已开源,我们诚挚邀请国内外的研究人员和开发者使用、复现和贡献代码,共同推动 RAG 技术的发展和国产化进程!