我要提问
ARTICLE DETAIL

资讯详情

前沿编程新知与开发实战干货的深度解读。

Agent Lightning Search-R1 实战:用 VERL 与 E5/FAISS 检索服务训练会“边思考边搜索”的问答 Agent

Agent Lightning Search-R1 实战:用 VERL 与 E5/FAISS 检索服务训练会“边思考边搜索”的问答 Agent Agent Lightning Search-R1 实战用 VERL 与 E5/FAISS 检索服务训练会“边思考边搜索”的问答 Agent【免费下载链接】agent-lightningThe absolute trainer to light up AI agents.项目地址: https://gitcode.com/GitHub_Trending/ag/agent-lightning本文围绕 Agent Lightning 仓库中的 Search-R1 示例examples/search_r1展开讲解如何用verl训练一个检索增强问答 Agent每个多轮 rollout 中模型在推理与 Wikipedia 检索之间交替进行最终输出答案并以精确匹配EM作为强化学习奖励。读完后你可以完成数据准备、独立部署 FAISS 稠密检索服务、通过run.sh一键启动agl-server 本地控制器 VERL 训练并在 Chat Completions 与 Token-in/token-out Completions 两种 API 模式之间切换。示例定位与运行前提该示例复现Search-R1: Training LLMs to Reason and Leverage Search Engines with Reinforcement LearningJin et al., 2025的核心思路官方参考配置如下见 docs/11-example-search-r1.mdGPUModelController ModeTrainer Mode8× A100 40GBmeta-llama/Llama-3.2-3B-InstructLocalSync only也就是说该示例使用 Agent Lightning 的本地控制器local controller以同步训练模式运行每条 rollout 都是一个本地进程通过 AGL Gateway 调用策略模型同时向一个独立的 FAISS 检索服务发起搜索请求。适用前提是机器具备约 8 张 A100 40GB 显存且已安装 Agent Lightning v1.0 及 conda 环境。整个示例由以下几部分组成retrieval_server.py基于 FastAPI FAISS E5 的稠密检索服务retrieval_launch.sh / data_process.sh检索服务启动与数据准备脚本agents/search_r1_agent.py多轮搜索 Agent 的核心实现含 Chat 与 Completion 两种 Agent 类agents/qa_em.pyEM 奖励计算train_search_r1_agent.py组装 VERL 配置并启动训练run.sh一键启动脚本。数据准备在仓库根目录执行examples/search_r1/data_process.sh该脚本会创建一个名为retriever的 Conda 环境Python 3.10、PyTorch 2.4.0 CUDA 12.1、faiss-gpu1.8.0、transformers4.57.1等见 data_process.sh并下载、解压以下产物examples/search_r1/data/wiki-18.jsonlWikipedia-18 语料从 Hugging Face 的PeterJinGo/wiki-18-corpus下载后gzip -dk解压examples/search_r1/data/e5_Flat.indexE5 稠密向量 FAISS 索引由part_aa、part_ab两个分片cat合并而成因为索引文件较大做了分片存储examples/search_r1/data/train.parquet与examples/search_r1/data/test.parquetNQ HotpotQA 问答训练/评测数据。脚本是幂等的对已存在且非空的文件直接skip existing重复执行不会重新下载。通过环境变量可以覆盖默认行为环境变量默认值说明SEARCH_R1_DATA_DIRexamples/search_r1/data语料与索引的存放目录SEARCH_R1_RETRIEVER_ENVretrieverConda 环境名SEARCH_R1_RETRIEVER_PYTHON3.10Conda 环境的 Python 版本SEARCH_R1_PYTORCH_CUDA12.1PyTorch 依赖的 CUDA 版本SEARCH_R1_SKIP_RETRIEVER_INSTALL0设为1时跳过检索依赖安装适合已有现成环境检索服务在另一个终端中启动检索服务并在整个训练期间保持其运行examples/search_r1/retrieval_launch.shretrieval_launch.sh 激活retrieverConda 环境后以默认参数调用retrieval_server.py并始终带--faiss-gpu将索引尝试迁移到 GPU。服务默认监听http://127.0.0.1:8000/retrieve可用健康检查端点确认就绪curl http://127.0.0.1:8000/healthz常用检索配置均为retrieval_launch.sh读取的环境变量配置项默认值说明SEARCH_R1_DATA_DIRexamples/search_r1/data语料与 FAISS 索引目录SEARCH_R1_RETRIEVAL_PORT8000检索服务端口SEARCH_R1_TOPK3每次搜索返回的文档数SEARCH_R1_RETRIEVER_DEVICEauto编码器设备auto、cuda、cuda:0或cpu从源码看retrieval_server.py 的检索链路是POST /retrieve接收QueryRequest字段queries、可选topk、return_scores→DenseEncoder.encode用intfloat/e5-base-v2对查询编码E5 家族模型会自动给每条查询加query:前缀并做 mean-pooling L2 归一化→ FAISS 索引search批量取 topk → 从语料取出title/text/contents与相似度分数返回。若安装的 faiss 没有 GPU 辅助函数或 CUDA 不可用源码会降级回 CPU 索引并打印警告而不是直接失败。此外retrieval_server.py还支持--query-max-length默认 256、--batch-size默认 512、--use-fp16等直接传给 uvicorn 服务的命令行参数便于在显存紧张时调节吞吐。训练启动在检索服务运行的前提下从仓库根目录执行examples/search_r1/run.shrun.sh 的完整流程是注册cleanup陷阱在脚本退出/中断时pkill掉agl-server、agl-controller进程并ray stop --force保证不残留后台进程将仓库根目录加入PYTHONPATHAgent 类以examples.search_r1...模块路径被控制器导入因此必须从仓库根启动后台启动agl-server端口默认8080key 默认dummy并指定default_proxy.model_name为策略模型随后轮询/healthz最多 60 秒等待就绪后台启动本地agl-controllerrunner_typelocal指向刚启动的 server前台运行python examples/search_r1/train_search_r1_agent.py把--agl-base-url、--agl-key、--run-name local以及用户追加的透传参数一并交给训练脚本。两种 API 模式默认 Agent 使用Chat Completions API文本进/文本出。若要改用Token-in/token-out Completions API——Agent 直接发送 prompt token ID 并接收生成 token ID从而在多轮过程中完整保留 token 序列——只需透传参数examples/search_r1/run.sh --api-type completion要指定不同的数据集文件examples/search_r1/run.sh \ --train-file /path/to/train.parquet \ --val-file /path/to/test.parquetAgent 运行时的配置项由 rollout 进程以环境变量读取配置项默认值说明SEARCH_R1_RETRIEVAL_URLhttp://127.0.0.1:8000/retrieverollout Agent 使用的检索端点SEARCH_R1_MAX_TURNS4每次 rollout 的最大模型/搜索轮数SEARCH_R1_MAX_TOKENS500每次模型回复的最大生成 token 数SEARCH_R1_TEMPERATURE1.0采样温度--api-type的切换逻辑在 train_search_r1_agent.py 的build_config中chat时控制器使用examples.search_r1.agents.search_r1_agent:SearchR1Agentcompletion时替换为同模块下的SearchR1CompletionAgent并通过env_map注入SEARCH_R1_TOKENIZER_MODEL取值就是当前策略模型路径供 Agent 做 prompt 编码与生成结果解码。Agent 内部机制多轮“推理—搜索”循环SearchR1Agent.run 是 rollout 的核心。它从本地控制器注入的环境变量读取任务与连接信息QUESTION问题、GOLDEN_ANSWERS标准答案支持 JSON 列表、AGL_KEY/AGL_EVENT_URL奖励上报以及AGL_OPENAI_BASE_URL指向 AGL Gateway 的 OpenAI 兼容端点。每一轮的循环逻辑search_r1_agent.py通过 Gateway 调用chat.completions.create获取模型回复postprocess_response在第一个/search或/answer处截断回复避免模型“越界”输出extract_action用正则(search|answer)(.*?)/\1解析动作。若解析失败标记invalid_action并终止本轮 rollout奖励记 0若动作是searchretrieve_doc以{queries: [query], topk: topk, return_scores: true}请求检索服务把命中文档格式化为Doc 1(Title: ...) ...文本并包裹进information.../information作为下一轮 user 消息若动作是answerexecute_response返回空字符串rollout 正常结束若max_turns耗尽且未结束会额外再做一次模型调用给模型最后一次输出答案的机会。SearchR1CompletionAgent的循环结构完全一致差别在于上下文维护方式Chat 版维护messages对话列表Completion 版则维护纯 token 序列——prompt_ids response_ids feedback_ids逐轮拼接search_r1_agent.py检索反馈也用同一个 tokenizer 编码后直接追加到 token 序列中从而避免文本往返带来的 token 不一致问题。提示词模板INSTRUCTION_FORMATsearch_r1_agent.py明确要求模型每次获得新信息后先在think标签内推理缺知识时用search query /search触发检索结果会出现在information.../information中可以搜索任意多次不再需要外部知识时直接给出answer.../answer。该模板注释中说明其改编自 Search-R1 原始仓库的数据处理脚本。奖励计算基于 EM 的稀疏奖励rollout 结束后Agent 会调用 qa_em.py 的compute_score_em计算奖励并通过 AGL 事件接口上报event_type: rewardsource: agentextract_solution抽取最后一个answer.../answer内容normalize_answer做小写化、去标点、去冠词a/an/the、空白归一化后与标准答案列表做精确匹配em_check命中得 1.0格式正确但未命中得 0.0若整段轨迹中从未出现合法动作invalid_action奖励直接记 0reason 为invalid_action否则 reason 为em_match或em_miss见 search_r1_agent.py。VERL 训练配置要点train_search_r1_agent.py 中的verl_default_config在 Agent Lightning 的基础 VERL 配置agentlightning/verl/config.yaml之上叠加了一批面向 Search-R1 的覆盖项关键的几组算法algorithm.adv_estimator: grpouse_kl_in_reward: Falseactor 侧开启 KL 惩罚use_kl_loss: Truekl_loss_type: low_var_klkl_loss_coef: 0.001数据与序列长度train_batch_size: 512max_prompt_length: 6000max_response_length: 4096truncation: errorRollout 引擎vLLMtensor_model_parallel_size: 1每条样本采样n: 4个轨迹配合 GRPO 组内相对优势gpu_memory_utilization: 0.5max_model_len: 32768多轮格式llama3_json并启用enable_auto_tool_choiceActor 优化ppo_mini_batch_size: 256、ppo_micro_batch_size_per_gpu: 4、学习率1e-6、clip_ratio_low/high: 0.2/0.3FSDP 开启param_offload与optimizer_offload以节省显存训练调度n_gpus_per_node: 8、total_epochs: 15、total_training_steps: 300、test_freq: 10、save_freq: 10日志输出到 console 与 wandbcheckpoint 目录为checkpoints/search_r1_checkpoints/Agent Lightning 接入train_search_r1_agent.pyagl_base_url: http://localhost:8080、rollout_timeout_seconds: 1800local.agent_class指定 Agent 类env_map将数据集字段input.question、input.golden_answers映射为 Agent 环境变量QUESTION、GOLDEN_ANSWERS。trace_aggregator采用trajectory级别聚合最大 prompt/response 长度分别为 4096 与 34384用于把多轮搜索轨迹收敛成一条训练轨迹。build_config使用 Hydra/OmegaConf 将基础配置、上述覆盖项以及命令行透传的dotlist覆盖项三者合并最终由run_ppoagentlightning/verl/entrypoint.py拉起 VERL 训练训练脚本还会在启动前把解析后的完整配置打印出来 VERL Config 方便核对。run.sh中透传的--run-name local会让实验名变为search_r1_local。相关文件索引文件作用examples/search_r1/README.md示例入口指向官方文档 docs/11-example-search-r1.mdexamples/search_r1/data_process.sh数据下载、索引合并与 Conda 环境准备examples/search_r1/retrieval_server.pyE5 FAISS 稠密检索 FastAPI 服务examples/search_r1/retrieval_launch.sh检索服务启动脚本examples/search_r1/agents/search_r1_agent.pySearchR1Agent与SearchR1CompletionAgent多轮 rollout 实现examples/search_r1/agents/qa_em.pyEM/SubEM 奖励计算examples/search_r1/train_search_r1_agent.pyVERL 配置组装与训练入口examples/search_r1/run.sh一键启动 server、本地控制器与训练适用限制该示例默认面向 8× A100 40GB 单机、Llama-3.2-3B-Instruct、本地控制器同步训练模式数据准备依赖 conda 与外部数据源的可达性检索服务需在整个训练期间常驻--api-type completion模式下 Agent 使用与策略模型相同的 tokenizer 做编解码因此SEARCH_R1_TOKENIZER_MODEL或由SEARCH_R1_MODEL回退需指向与模型匹配的 tokenizer。【免费下载链接】agent-lightningThe absolute trainer to light up AI agents.项目地址: https://gitcode.com/GitHub_Trending/ag/agent-lightning创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表