企业官网建设流程全解析

1. 项目概述：OneGen，让大模型“一次生成，顺便检索”

如果你正在折腾大语言模型（LLM）的应用，尤其是检索增强生成（RAG）这个方向，那你肯定对“先检索，再生成”这个标准流程又爱又恨。爱的是，它确实能让模型回答得更准、更靠谱；恨的是，这流程太“笨重”了——你得先跑一遍检索模型（比如一个向量数据库查询），拿到相关文档，再把这些文档塞给生成模型去处理。这相当于让模型做了两次“前向传播”，计算开销和延迟都上去了，部署成本也水涨船高。

最近我在跟进一个叫OneGen的开源项目，它来自浙江大学 ZJUNLP 实验室，思路非常巧妙。它的核心目标就写在名字里：One-Pass Unified Generation and Retrieval，即“一次前向传播，统一完成生成和检索”。简单来说，它训练一个大模型，让它自己学会在生成答案的过程中，顺便把该检索的“检索令牌”也一并生成了，从而把检索动作内化到生成流程里。这听起来有点像让模型“一心二用”，但实现上却非常优雅，直接省掉了一次独立的检索模型前向计算，还能复用生成过程中的 KV Cache，推理效率提升显著。

我花了一些时间把它的代码、论文和实验都捋了一遍，发现这不仅仅是一个学术玩具，其设计理念和工程实现对于想构建高效、低成本 RAG 系统的开发者来说，有很强的参考价值。它目前主要支持三个任务：实体链接、单跳问答和多跳问答，基于 Llama2-7B 等模型进行了微调，所有模型和代码都已开源。接下来，我就结合自己的理解，带你深入拆解 OneGen 的核心思想、实操细节，并分享在复现和尝试过程中遇到的一些坑和心得。

2. 核心设计思路：令牌的“角色扮演”与统一训练

要理解 OneGen，首先得跳出“模型要么在生成，要么在检索”的二分法思维。OneGen 提出了一个更底层的视角：从令牌（Token）的角色出发。

2.1 令牌的三种角色

在 OneGen 的框架里，输入给模型的每一个令牌x_i都被赋予了一个明确的“角色”：

• 生成角色（GEN）：这个令牌的目标是参与预测下一个令牌。这是标准语言模型训练的核心，计算交叉熵损失。
• 上下文角色（CTX）：这个令牌的作用是为后续的生成提供背景信息。它本身不直接参与损失计算，但会影响模型对后续令牌的预测。
• 检索角色（RET）：这是 OneGen 的创新关键。这个令牌代表了一个需要被检索的“句子”或“片段”的语义。它的训练目标不是生成文本，而是让模型学会输出一个能代表某段文本语义的“向量表示”。

举个例子，在一个多跳问答任务中，用户问：“爱因斯坦在哪个大学获得了博士学位？这个大学所在的城市以什么闻名？” 模型在生成过程中，可能会先输出一个 RET 令牌，这个令牌的隐含向量（即模型最后一层的隐藏状态）应该与知识库中“苏黎世联邦理工学院”这个实体的语义向量尽可能接近。然后，模型再基于这个“检索到”的语义，继续生成后续的答案。

2.2 统一的训练目标

基于令牌角色，OneGen 设计了一个混合损失函数：

• 对于GEN角色的令牌，使用标准的交叉熵损失，确保模型的语言生成能力。
• 对于RET角色的令牌，使用对比损失。具体来说，模型会为每个 RET 令牌生成一个向量表示，这个表示要与对应的“正样本”句子向量（比如从知识库中提取的正确实体向量）在向量空间里尽可能接近，同时与一批“负样本”句子向量尽可能远离。

这样，在一个训练批次中，模型同时在学习两件事：1）如何流畅地生成文本；2）如何输出有意义的向量来表示特定语义内容，以便进行检索。这才是“统一”的精髓——不是两个模型拼接，而是一个模型同时掌握两种技能。

2.3 与主流方案的对比

为了更直观地理解 OneGen 的优势，我们可以把它和几种常见的 RAG 实现方式做个对比：

从对比可以看出，OneGen 在理论上实现了最精简的推理路径。它避免了为检索单独对查询进行编码，直接利用生成过程中的中间状态来完成检索匹配，这对降低端到端延迟和计算成本非常有利。

3. 环境搭建与数据准备实操

理论很美好，但能不能跑起来才是关键。OneGen 的代码库结构比较清晰，但有些细节需要特别注意。

3.1 基础环境配置

项目推荐使用 Python 3.9 和 Conda 环境，这是为了避免一些潜在的包依赖冲突。

# 1. 克隆代码库 git clone https://github.com/zjunlp/OneGen cd OneGen # 2. 创建并激活 Conda 环境 conda create -n onegen python=3.9 -y conda activate onegen # 3. 安装依赖 pip install -r requirements.txt

这里有个小坑：requirements.txt里固定了某些库的版本（比如transformers）。如果你本地已经有其他深度学习项目，可能会产生冲突。我的建议是专门为 OneGen 创建这个独立环境，不要和其他项目混用。安装过程如果遇到网络问题，可以考虑为 pip 配置镜像源。

3.2 数据获取与处理

论文实验使用了三个任务的数据。项目提供了两种数据获取方式：

1. 从 Google Drive 下载处理好的数据（主要用于评估）：

   • 下载eval_data.tar.gz。
   • 解压后得到eval_data文件夹，将其放入项目根目录的data/目录下。

   tar -xzvf eval_data.tar.gz mv eval_data ./data/

   • 注意：train_data.tar.gz作者说可以从 Hugging Face 加载，所以不一定需要下载。

2. 直接从 Hugging Face Datasets 加载（用于训练）：

   • 这是更推荐的方式。训练脚本（train.py）内部会通过 Hugging Face 的datasets库下载对应任务的数据集。你需要确保网络能够访问https://huggingface.co。
   • 如果遇到下载慢或失败的情况，可以尝试设置环境变量HF_ENDPOINT=https://hf-mirror.com来使用镜像站。

重要提示：对于**实体链接（Entity Linking）**任务，评估时需要用到预计算好的实体向量（Embeddings）。这个文件（OneGen-EntityLinking-Llama2-7B-Embedding.pkl）必须单独从 Hugging Face 仓库下载，并放在正确路径下，否则评估脚本会报错。这是新手最容易忽略的一点。

3.3 模型下载

OneGen 提供了在三个任务上微调好的 Llama2-7B 模型，发布在 Hugging Face、ModelScope 和 WiseModel 平台。我们可以直接用git-lfs克隆到本地。

# 例如，下载实体链接模型 git lfs install git clone https://huggingface.co/zjunlp/OneGen-EntityLinking-Llama2-7B

由于模型较大（约14GB），请确保本地有足够的磁盘空间和稳定的网络。如果使用 ModelScope，可以使用其提供的客户端工具modelscope进行下载。

4. 模型训练与微调详解

如果你想在自己的数据集上应用 OneGen 的思想，或者复现论文结果，就需要进行训练。项目使用 DeepSpeed 进行分布式训练，以应对 7B 模型参数量。

4.1 配置文件解析

训练的核心是配置文件，位于workflow/{task}/{model}.json。以workflow/entity_linking/llama2.json为例，我们看看几个关键参数：

{ “info_model”: { “model_path”: “meta-llama/Llama-2-7b-hf”， // 基础模型路径 “tokenizer_path”: “meta-llama/Llama-2-7b-hf”, “max_length”: 1024 // 序列最大长度 }, “train_args”: { “output_dir”: “./output/el_llama2”， “num_train_epochs”: 3， “per_device_train_batch_size”: 4， // 每张 GPU 的批大小 “gradient_accumulation_steps”: 4， // 梯度累积步数 “learning_rate”: 2e-5， “deepspeed”: “./configs/ds_config.json” // DeepSpeed 配置文件 }, “data_args”: { “n_pos_per_sent”: 5， // 每个 RET 令牌对应的正样本数 “n_neg_per_pos”: 7， // 每个正样本对应的负样本数 “dataset_name”: “zjunlp/OneGen-EntityLinking” // 数据集名称 } }

• n_pos_per_sent和n_neg_per_pos：这是对比学习的关键。对于每一个 RET 令牌，需要为其构造正样本和负样本。这些参数控制了样本数量，直接影响训练效果和内存占用。值越大，对比学习任务越难，效果可能更好，但显存消耗也越大。
• per_device_train_batch_size：在 8x A800 (80GB) 上设置为 4。如果你用的是显存较小的卡（如 24GB 的 3090/4090），必须调小这个值，比如设为 1，同时可能需要调整梯度累积步数以保持总的有效批次大小。
• deepspeed：项目提供了 DeepSpeed 的 ZeRO-2 配置文件，支持优化器状态和梯度的分片，是能训练 7B 模型的关键。除非你非常熟悉 DeepSpeed，否则不建议修改这个配置文件。

4.2 启动训练

假设你已经配置好 8 卡 A800 环境，并且数据已就绪，启动训练非常简单：

# 实体链接任务 deepspeed train.py --workflow workflow/entity_linking/llama2.json # 单跳问答任务 deepspeed train.py --workflow workflow/self_rag/llama2.json # 多跳问答任务 deepspeed train.py --workflow workflow/multi_hop_qa/llama2.json

实操心得：

1. 监控显存：训练开始后，立刻用nvidia-smi观察显存占用。如果看到显存迅速被占满然后报 OOM（内存不足），首要任务就是减小per_device_train_batch_size。
2. 学习率预热：DeepSpeed 配置中通常包含了学习率调度。如果训练损失一开始就爆炸（变成 NaN），可能是学习率太高，可以尝试在train_args中微调learning_rate，例如从2e-5降到1e-5。
3. 日志与保存：训练日志和模型检查点会保存在output_dir指定的目录。定期检查损失曲线，确保生成损失和检索对比损失都在平稳下降。

4.3 单卡或资源不足的应对策略

如果你只有单卡或显存不足，直接运行上述命令肯定会失败。你需要一个组合拳：

1. 降低批次大小：将per_device_train_batch_size设为 1。
2. 启用梯度检查点：在info_model部分添加“gradient_checkpointing”: true。这会用计算时间换显存，大约能节省 20%-30% 的显存。
3. 使用更低精度：修改./configs/ds_config.json，将“fp16”: {“enabled”: true}改为使用“bf16”: {“enabled”: true}（如果你的 GPU 支持 BF16，如 A100/A800），或者尝试“fp16”: {“enabled”: true, “loss_scale”: 0, “loss_scale_window”: 1000, “initial_scale_power”: 16}并进行梯度缩放管理。更激进的方法是使用“fp8”: {“enabled”: true}（H100 支持）。
4. 使用参数卸载：在 DeepSpeed 配置中启用 ZeRO-3，并将优化器状态、梯度甚至模型参数卸载到 CPU 内存。但这会显著增加 CPU-GPU 通信，大幅降低训练速度。

对于大多数拥有 24GB 显存的用户，采用“批次大小=1 + 梯度检查点 + BF16/FP16”的策略，有希望跑起来 7B 模型的训练，但每一步的时间会很长。另一种思路是考虑使用QLoRA等参数高效微调方法，但 OneGen 原代码目前并未集成，需要自己动手修改适配，这是一个进阶挑战。

5. 模型推理与评估实战

训练好（或下载好）模型后，下一步就是用它来做预测和评估。OneGen 的推理过程巧妙地体现了其“单次前向”的优势。

5.1 推理脚本配置与运行

推理入口是eval.py，它需要一个 JSON 配置文件。以实体链接为例，配置文件是config/eval_config/entity_linking/llama2_wo_pkl.json。

{ “model”: { “model_path”: “zjunlp/OneGen-EntityLinking-Llama2-7B”， “tokenizer_path”: “zjunlp/OneGen-EntityLinking-Llama2-7B”， “max_length”: 1024 }， “data”: { “file”: “./data/eval_data/entity_linking/test.json”， // 输入数据文件 “output_file_path”: “./outputs/el_results.jsonl” // 输出结果文件 }， “inference”: { “use_faiss”: true， // 是否使用 Faiss 加速检索 “index_path”: “./data/eval_data/entity_linking/entity_index.faiss”， // Faiss 索引文件 “embedding_path”: “./data/eval_data/entity_linking/entity_embedding.pkl” // 实体向量文件 } }

关键配置项：

• use_faiss：强烈建议设置为true。Faiss 是 Facebook 开源的向量相似性搜索库，对于在大规模实体库（数万甚至数百万）中查找与 RET 令牌向量最接近的实体，它能提供成百上千倍的加速。如果设为false，则会使用纯 Python 循环计算余弦相似度，在实体数量稍多时就会慢得无法忍受。
• index_path和embedding_path：这两个文件通常包含在之前下载的eval_data中。确保路径正确。

运行推理：

# 实体链接推理 python eval.py --config config/eval_config/entity_linking/llama2_wo_pkl.json # 多跳问答推理 python eval.py --config config/eval_config/multi_hop_qa/llama2.json

推理过程发生了什么？

1. 模型加载输入（例如一个问题）。
2. 模型开始自回归生成。当它遇到需要检索的位置时（由训练决定），它会输出一个特殊的 RET 令牌。
3. 程序捕获这个 RET 令牌对应的隐藏状态向量。
4. 用这个向量在 Faiss 索引中搜索最相似的实体向量（对于实体链接）或文档向量（对于问答）。
5. 将检索到的实体/文档信息以某种方式（例如，直接插入文本）反馈给模型上下文，模型继续生成后续内容。
6. 循环直至生成结束。

这个过程是“单次前向”的，因为检索动作发生在生成流内部，模型没有停下来等待一个外部的检索调用。

5.2 结果评估

推理完成后，会生成一个.jsonl文件，每行包含模型对一条输入数据的预测结果。接下来需要用专门的脚本来计算指标。

# 评估实体链接结果 bash scripts/eval_el.sh el /path/to/your/result.jsonl # 评估多跳问答结果 (HotpotQA 数据集) bash scripts/eval_multi_hop_qa.sh /path/to/your/result.jsonl hotpotqa # 评估多跳问答结果 (2Wiki 数据集) bash scripts/eval_multi_hop_qa.sh /path/to/your/result.jsonl 2wiki

这些评估脚本大多是 Python 写的，会计算诸如准确率（Accuracy）、F1 值、精确匹配（EM）等任务相关指标。运行它们通常只需要 CPU。

对于单跳问答（Self-RAG），评估方式比较特殊，因为它集成了 Self-RAG 的评估框架，需要调用模型本身进行检索和生成决策：

CUDA_VISIBLE_DEVICES=0 bash scripts/eval_self_rag.sh 0 always_retrieve /path/to/model model_tag saved_rank_path 5 true true

这个脚本参数较多，核心是always_retrieve模式，即强制模型在每一个决策点都进行检索。你需要确保saved_rank_path下有检索库的相关数据。

5.3 常见问题与排查

在跑通整个流程时，我遇到了几个典型问题，这里分享解决方案：

1. 报错：KeyError: ‘retrieval_token_id’

   • 原因：模型的config.json文件里缺少 OneGen 自定义的特殊令牌 ID 定义。
   • 解决：检查你使用的模型路径。如果是下载的官方模型，不应该有此问题。如果是自己训练的，需要在保存模型时，确保tokenizer和model.config都正确添加了特殊令牌（如<ret>）。可以在训练配置的info_model里添加“special_tokens”: {“retrieval_token”: “<ret>”}等相关设置，并确保代码正确处理了这些令牌的添加。

2. 评估实体链接时找不到.pkl文件

   • 原因：没有下载预计算的实体向量文件。
   • 解决：从 Hugging Face 仓库zjunlp/OneGenEmbedding下载OneGen-EntityLinking-Llama2-7B-Embedding.pkl，并确保它在inference.embedding_path指定的路径上。

3. Faiss 索引加载失败

   • 原因：Faiss 索引文件损坏，或者是用不同版本的 Faiss 创建的。
   • 解决：重新从eval_data中解压索引文件。如果问题依旧，尝试重新生成索引（如果你有原始的实体向量）。通常直接使用项目提供的数据包是最稳妥的。

4. 推理速度慢

   • 检查点：首先确认use_faiss是true。其次，检查生成的序列长度是否过长（max_length设置是否过大）。对于批处理推理，可以尝试在配置中增加batch_size参数（如果代码支持）。

5. 显存不足（OOM） during inference

   • 原因：即使推理比训练省显存，但处理长序列或大模型时，7B 模型在单卡上也可能 OOM。
   • 解决：减少max_length；启用模型本身的device_map=“auto”或load_in_8bit/load_in_4bit（需要修改代码以支持 bitsandbytes 量化加载）。

6. 深入原理：RET令牌如何工作与训练技巧

看完了实操，我们回过头再深入一点，聊聊 OneGen 里最核心的“魔法”——RET 令牌到底是怎么被训练和使用的。这部分理解透了，你才能更好地把它应用到自己的任务上。

6.1 RET令牌的训练数据构造

模型怎么知道哪个位置该输出 RET 令牌？这完全依赖于训练数据的标注。在准备训练数据时，你需要：

1. 确定检索点：在文本序列中，人工或通过启发式规则确定哪些位置需要进行检索。例如，在多跳问答中，第一个问题答案之后、第二个问题开始之前，可能就需要一个检索点，去获取回答第二个问题所需的新知识。
2. 插入特殊令牌：在检索点插入一个特殊的、代表“检索”的令牌，比如<ret>。这个令牌在词汇表中是唯一的。
3. 准备正负样本：对于每一个<ret>令牌，你需要知道它“应该”检索到什么。这个“应该检索到的文本”的向量，就是正样本。同时，你需要从知识库中采样一些其他文本的向量作为负样本。

在代码中，DataCollator会负责在构建批次时，根据这些标注，为每个 RET 令牌组装好对应的正样本向量和负样本向量列表，然后传递给模型计算对比损失。

6.2 对比损失的具体实现

OneGen 使用的对比损失通常是 InfoNCE Loss 的一种变体。对于一个 RET 令牌的向量v，其正样本向量为p，负样本向量集合为{n1, n2, ..., nk}，损失函数鼓励v与p的相似度远高于与所有n的相似度。

sim()通常是余弦相似度。这个损失函数会与 GEN 令牌的交叉熵损失加权求和，共同指导模型参数更新。权重的选择是一个超参数，需要小心调整，以平衡生成质量和检索能力。

6.3 推理时的检索-生成协同

推理时，模型以自回归方式运行。当它生成到<ret>这个令牌时：

1. 前向传播计算出的、对应<ret>位置的最后一层隐藏状态被取出，作为查询向量q。
2. 用q在 Faiss 索引中搜索最近邻，找到最相关的文本片段d。
3. 接下来是关键：如何把d的信息给回模型？论文和代码中可能采用了几种策略：
   • 直接拼接：将d的文本直接拼接到已生成的序列后面，作为后续生成的上下文。这简单，但可能破坏连贯性。
   • 软性注入：将d的向量表示以某种方式（如注意力）融合到后续的生成过程中。这更复杂，但可能更流畅。
   • 标记化后输入：将d文本 tokenize 后，作为后续的输入序列。OneGen 可能采用了类似这种方式，因为它保持了纯粹的自回归生成框架。

你需要仔细阅读eval.py中生成循环部分的代码，才能确定具体策略。理解这一点对于你后续自定义任务至关重要。

7. 扩展思考与应用展望

OneGen 提供了一个非常漂亮的框架原型，但要把它的思想用到自己的生产环境中，还需要考虑很多工程和算法上的扩展。

7.1 支持更多任务类型

目前官方支持了三个任务。你可以尝试将它扩展到：

• 长文本摘要：在生成摘要的不同部分时，动态检索原文的相关段落作为依据，避免遗忘。
• 代码生成：生成代码时，检索相似的 API 使用范例或代码片段。
• 对话系统：在多轮对话中，根据当前对话历史检索相关知识库条目，再生成回复。

关键在于如何定义你的“检索点”和“检索单元”。对于对话，检索点可能在每轮回复开始前；对于代码生成，可能在生成每个函数或类之前。

7.2 与现有RAG系统集成

你不需要从头训练一个 OneGen 模型。一个折中的思路是：使用一个训练好的 OneGen 模型作为“检索决策器”。

1. 让 OneGen 模型先处理用户查询，它会在内部生成一些 RET 令牌及其对应的向量。
2. 将这些向量用于你的现有向量数据库进行检索。
3. 将检索结果和原始查询一起，交给另一个更强大的、专精生成的模型（如 GPT-4、Claude 3）来生成最终答案。

这样，你利用了 OneGen 高效、内生的检索能力来改善检索质量，同时保留了顶级生成模型的强大能力。

7.3 工程化改进方向

项目目前的 TODO list 也指出了几个有价值的改进方向：

• 支持 LoRA 训练：这是降低训练成本、实现轻量化微调的必经之路。可以基于peft库进行集成。
• 支持 vLLM 推理：vLLM 的 PagedAttention 能极大提升大模型推理的吞吐量。将 OneGen 与 vLLM 结合，能充分发挥其高效单次前向的优势。
• 分布式嵌入：当知识库极大时，单个 Faiss 索引可能放不进内存。需要支持分布式向量索引，如 Faiss 的IndexShards。
• Gradio 演示界面：构建一个 Web 界面，直观展示模型如何一步步生成并检索，对于理解和演示项目价值巨大。

7.4 潜在挑战与注意事项

1. 训练稳定性：联合训练生成和对比损失并不容易。两种损失的量级和收敛速度可能不同，需要仔细调整损失权重、学习率调度和批次采样策略。
2. 检索质量依赖：模型检索的好坏，极度依赖训练时构造的正负样本质量。负样本如果太简单，模型学不到区分能力；如果太难或噪声大，又会干扰学习。
3. 错误传播：在推理链中，如果前一步的 RET 检索错了，会直接导致后续生成基于错误信息，产生“一本正经胡说八道”的结果。需要考虑如何引入一定的验证或回溯机制。
4. 领域适配：在一个领域（如维基百科问答）上训练的 RET 令牌向量，可能无法直接迁移到另一个领域（如医疗文献）。需要进行领域适配微调。

折腾完 OneGen 的代码和实验，我最深的体会是，它代表了一种让大模型变得更“自主”和“高效”的思路。它不再被动地等待外部检索的结果，而是主动在思维流中标记出需要外部知识的节点，并自己动手去获取。这种“生成式检索”的范式，或许会成为未来构建复杂、可靠 AI 系统的一个重要组件。虽然现在它还有不少限制，但开源代码和模型为我们提供了一个绝佳的起点，剩下的，就是结合我们自己的具体问题，去迭代、优化和创造了。如果你正在为 RAG 系统的延迟和成本发愁，OneGen 绝对值得你花时间深入探究一番。