字节 Seed 团队联合清华大学智能产业研究院开源 MemAgent: 基于多轮对话强化学习记忆代理的长文本大语言模型重构

发布于：2025-07-11 ⋅ 阅读:(97) ⋅ 点赞:(0)

在这里插入图片描述

🔥 最新动态!!!

  • [2025/07] 我们提供了快速启动脚本，让使用MemAgent变得超级简单，详情请见下方"快速入门"部分。
  • [2025/06] 我们发布了RL-MemAgent-14B和RL-MemAgent-7B模型，在350万token上下文任务中实现了近乎无损的性能表现。
    在这里插入图片描述

在这里插入图片描述

📖 简介

我们提出了一种新颖的长上下文处理框架——MemAgent，它无需改变基础模型架构，直接通过端到端的强化学习优化长上下文任务。MemAgent 展现出卓越的长上下文处理能力：在 32K 文本上训练后，可外推至 3.5M 问答任务且性能损失 <5%，并在 512K RULER 测试中达到 95%+ 准确率。

在这里插入图片描述

亮点：

  • 🚀 创新记忆机制 引入MemAgent架构，可在固定上下文窗口内处理任意长度的输入，突破传统上下文窗口长度限制。
  • ⚡ 线性时间复杂度 突破长文本处理的计算瓶颈，实现资源消耗与文本长度的线性扩展。
  • 🎯 强化学习驱动的外推能力 通过MemAgent架构的强化学习训练，使模型能够外推到更长的文本而性能损失最小。

在这里插入图片描述

多轮对话强化学习框架

我们采用可验证奖励的强化学习（RLVR）来训练记忆代理，扩展了DAPO算法以支持具备多轮上下文无关对话的智能体工作流程端到端优化。

在这里插入图片描述

结果

RL-MemAgent 在超长上下文处理中展现出卓越的稳定性：

  • 140亿参数模型：在350万token的任务上性能衰减<5.5%，实现真正无损的外推。
  • 70亿参数模型：在最长上下文中性能仅下降11%，显著优于现有长上下文模型。

在这里插入图片描述

快速开始

quickstart.py 提供了一种简单上手 MemAgent 的方式，支持本地部署和在线模型服务集成。

vLLM 本地部署

  1. 启动 vllm 服务器:

    vllm serve BytedTsinghua-SIA/RL-MemoryAgent-14B --tensor_parallel_size 2

  2. 运行 quickstart.py:

    python quickstart.py --model BytedTsinghua-SIA/RL-MemoryAgent-14B

在线大语言模型服务

对于在线大语言模型服务，您需要将模型端点地址和API密钥配置为环境变量。

例如 gpt-4o-2024-11-20：

  • 普通在线服务：直接使用 https://{endpoint}。
  • Azure OpenAI：使用格式 https://{endpoint}/openai/deployments/gpt-4o-2024-11-20。
export URL=
export API_KEY=
python quickstart.py --model gpt-4o-2024-11-20

可复现性

测试结果

pip install httpx==0.23.1 aiohttp -U ray[serve,default] vllm
  1. 准备问答数据
cd taskutils/memory_data
bash download_qa_dataset.sh
  1. 下载数据集
cd ../..
bash hfd.sh BytedTsinghua-SIA/hotpotqa --dataset --tool aria2c -x 10
export DATAROOT=$(pwd)/hotpotqa
  1. 准备模型

测试中使用的模型将从HuggingFace下载。但Qwen2.5-Instruct系列模型需要手动下载，并正确配置其config.json文件以启用YaRN功能。请按照Qwen2.5-Instruct仓库的说明进行操作。

bash hfd.sh Qwen/Qwen2.5-7B-Instruct --model --tool aria2c -x 10
bash hfd.sh Qwen/Qwen2.5-14B-Instruct --model --tool aria2c -x 10
bash hfd.sh Qwen/Qwen2.5-32B-Instruct --model --tool aria2c -x 10
# then change the config.json manually

export MODELROOT=/your/path/to/models # move to your model root directory, this env variable is used in the run.py script
mv Qwen2.5-7B-Instruct $MODELROOT/Qwen2.5-7B-Instruct-128K
mv Qwen2.5-14B-Instruct $MODELROOT/Qwen2.5-14B-Instruct-128K
mv Qwen2.5-32B-Instruct $MODELROOT/Qwen2.5-32B-Instruct-128K
  1. 运行

注意： 运行所有测试需要几天时间，您可能需要指定要运行的测试/模型。

cd taskutils/memory_eval
python run.py

注意： 此脚本将使用所有可用GPU来部署模型。若您拥有多个GPU节点，可以创建一个Ray集群并在其中一个集群节点上运行脚本。使用SERVE_PORT和DASH_PORT为Ray集群指定端口。

cd taskutils/memory_eval
SERVE_PORT=8000 DASH_PORT=8265 python run.py # port numbers here are default values, you may need to specify them as the serve/dashboard port in your ray cluster

训练

首先在run_memory_7B.sh和run_memory_14B.sh文件中指定PROJ_ROOT（用于保存检查点）和DATASET_ROOT（用于训练数据，应与测试时使用的路径相同）。

然后直接运行该脚本启动单节点训练，或者正确配置Ray集群并在集群节点之一上运行该脚本。

数据

本节操作请在taskutils/memory_data目录下执行以下命令。

cd taskutils/memory_data
pip install nltk pyyaml beautifulsoup4 html2text wonderwords tenacity fire
  1. 训练 & dev split: hotpotqa_train.parquet & hotpotqa_dev.parquet
  • 下载问答数据集和合成数据，如果在上一步已经下载过，请跳过此步骤：
bash download_qa_dataset.sh

python processing.py # Dataprocess, synthetic long context multihop-QA

  • 在本地主机8000端口部署Qwen-7B，在本地主机8001端口部署Qwen-7B-Instruct

  • 过滤

python filter.py -i hotpotqa_dev_process.parquet -o hotpotqa_dev_result --noresume
python filter.py -i hotpotqa_train_process.parquet -o hotpotqa_train_result --noresume
python3 filter2.py # Filtering out sample which can be answered correctly by LLM without any context:

2. Main task: eval_{50|100|200|…}.json

export DATAROOT="your_dir_to_hotpotqa_dev.parquet"
python convert_to_eval.py # Convert the `hotpotqa_dev` to `eval_200.json`
python different_docs_eval.py.py # Create eval dataset with different number of documents

3. OOD task: eval_{rulersubset}_{8192|16384|…}.json

export DATAROOT="your_dir_to_hotpotqa_dev.parquet"
python download_paulgraham_essay.py
bash download_qa_dataset.sh
bash ruler_data_prepare.sh

工程实现

同步模式：从工具调用到通用工作流

受Search-R1启发，我们实现了一个具有独立上下文的通用多会话工作流框架。因此上下文不再像原始工具调用那样仅限于所有先前对话的串联字符串。该框架使得以端到端方式优化多步智能体成为可能。记忆智能体是该框架的示例之一，展示了如何应用强化学习来优化智能体在多步工作流中的性能。

实现细节请参阅recurrent/impls/memory.py，接口设计参见recurrent/interface.py和recurrent/generation_manager.py文件。

异步模式：将智能体视为函数

基于服务器模式生成的实现，我们进一步开发了一个全新框架——允许用户将智能体实现为以openai-api风格调用大语言模型的函数，无需关注批处理张量操作（如分词、填充、状态追踪等），这些操作通常需要大量模板代码甚至状态机管理。

在此框架下，每个智能体都是一个函数，其返回值为一个或多个{"role":"", "content":""}字典构成的列表：

在这里插入图片描述

  • 香草模型（基础对话）：[{"role":"user","content":用户输入},{"role":"assistant""content":助手回复}]
  • 多轮工具调用模式：[{"role":"user","content":用户输入},{"role":"assistant","content":助手初始回复},{"role":"tool","content":工具观测结果},{"role":"assistant","content":助手最终回复}]
  • 上下文无关的多轮对话：[{"role":"user","content":对话1输入},{"role":"assistant","content":对话1回复}], [{"role":"user","content":对话2输入},{"role":"assistant","content":对话2回复}]

我们修改了chat_template以支持无需张量操作的工具响应掩码功能。

具体实现请参阅recurrent/impls/async_*.py文件中定义的agent类的rollout方法，框架设计详见recurrent/interface.py和recurrent/async_generation_manager.py文件。

RayActor进程池

计算密集型的奖励计算或工具调用可能会阻塞头节点向大语言模型发送生成请求。为提高CPU密集型任务的效率，我们在每个节点创建了一个运行进程池的RayActor，用于接收来自头节点的任务。

因此，CPU任务可以通过提交给ray actor异步执行，GPU任务（如LLM生成）同理。

详见verl/workers/reward_manager/thread.py文件。

致谢

我们感谢verl团队提供了灵活强大的基础设施。

我们感谢服务器模式生成的作者们所做的杰出工作，这为我们的异步智能体框架奠定了坚实基础。

代码

https://github.com/BytedTsinghua-SIA/MemAgent

网站公告

今日签到

点亮在社区的每一天
去签到

热门文章

最新发布