Xirui1208/memagent-readtwice-7b-step140

MemAgent ReadTwice 7B — step 140

基于 Qwen2.5-7B-Instruct 用 MemAgent（RL）训练的 long-context memory agent，采用 two-pass / read-twice 方法。本仓库为训练到 global_step_140 的 checkpoint（HF safetensors 格式，bf16）。

⚠️ 这是 two-pass 方法的 ckpt。评测必须走 two-pass 流程（先 skim 全文建 outline，再带 outline 顺序 update memory）。如果用普通 single-pass harness 跑，体现不出方法效果。复现脚本与 prompt 模板都在本仓库 eval/ 目录里。

评测结果

benchmark	setting	F1	EM
RULER-HQA (length=400, ~48K tok)	chunk=10000, two-pass (recurrent-boxed), vLLM fp8 TP4	67.8	50.78

文件结构

.                              # 根目录：标准 HF 权重，可直接 from_pretrained / vLLM 加载
├── config.json, *.safetensors, tokenizer.*
└── eval/                      # two-pass 评测配方（复现用）
    ├── eval_step140.sh        # 一键评测脚本（起 vLLM + 跑 RULER-HQA）
    ├── memory_eval/           # 评测引擎；two-pass 逻辑 + prompt 全在 utils/recurrent_boxed.py
    ├── data/eval_400.json     # RULER-HQA length=400 评测数据
    ├── requirements.txt
    └── NOTES.md               # 方法说明 + 历史实验结果

快速加载（标准 single-pass，仅验证权重可用）

from transformers import AutoModelForCausalLM, AutoTokenizer
m = AutoModelForCausalLM.from_pretrained("<repo_id>", torch_dtype="bfloat16", device_map="auto")
tok = AutoTokenizer.from_pretrained("<repo_id>")

复现 two-pass 评测

# 1) 下载本仓库（含 eval/）
hf download <repo_id> --local-dir ./memagent-readtwice-step140

cd ./memagent-readtwice-step140/eval
pip install -r requirements.txt   # 需要 vllm + 评测依赖

# 2) 起 vLLM（4 卡 TP4，fp8；按你的卡数调整）
CUDA_VISIBLE_DEVICES=0,1,2,3 vllm serve ../ \
  --served-model-name huggingface --enable-prefix-caching \
  --quantization fp8 --tensor-parallel-size 4 --port 8001 --max-model-len 32768

# 3) 跑 RULER-HQA（two-pass = --api recurrent-boxed）
cd memory_eval
DATAROOT=../data SERVE_PORT=8001 RECURRENT_CHUNK_SIZE=10000 RECURRENT_MAX_NEW=1024 \
  python ruler_hqa.py --length 400 --save_dir /tmp/eval_step140 --save_file x \
  --model huggingface --tokenizer ../../ --n_proc 256 --api recurrent-boxed --sampling 1 --force

输出末尾的 f1: / em: 即为指标。eval_step140.sh 把上面三步串好了，把里面的 HF/路径改成你的本地目录即可一键跑。

关键点

• two-pass 入口：评测命令的 --api recurrent-boxed，对应 memory_eval/utils/recurrent_boxed.py。
• chunk 大小：训练/评测用 RECURRENT_CHUNK_SIZE=10000（默认 5000，务必显式设成 10000 以对齐本 ckpt）。
• prefix cache：chunk 放在每个 prompt 最前面，Pass 1 / Pass 2 之间复用 vLLM prefix cache，开 --enable-prefix-caching。
• prompt 模板（skim / update / final）全部内嵌在 recurrent_boxed.py，无需额外训练侧文件。

base model: Qwen/Qwen2.5-7B-Instruct