🔥 最新動態!!!

• [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

2. 下載數據集

cd ../..
bash hfd.sh BytedTsinghua-SIA/hotpotqa --dataset --tool aria2c -x 10
export DATAROOT=$(pwd)/hotpotqa

3. 準備模型

測試中使用的模型將從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

4. 運行

注意： 運行所有測試需要幾天時間，您可能需要指定要運行的測試/模型。

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