LongMem

我们论文《通过长期记忆增强语言模型》（arXiv:2306.07174）的官方实现。

如果您觉得本仓库有趣或有帮助，请引用我们的论文：

@article{LongMem,
  title={Augmenting Language Models with Long-Term Memory},
  author={Wang, Weizhi and Dong, Li and Cheng, Hao and Liu, Xiaodong and Yan, Xifeng and Gao, Jianfeng and Wei, Furu},
  journal={arXiv preprint arXiv:2306.07174},
  year={2023}
}

环境配置

• PyTorch：请按照 PyTorch 官方安装指南 进行安装。建议使用 PyTorch ≥ 1.8.0。请根据您的 CUDA 驱动版本选择对应的 GPU 版本。

• Faiss-GPU：对于 Nvidia V100 显卡，直接运行 pip install faiss-gpu 即可。对于 Nvidia A100 和 A6000 显卡，请运行 conda install faiss-gpu cudatoolkit=11.0 -c pytorch。需要注意的是，Faiss-GPU 尚未正式支持 A100 显卡，有时可能会出现错误，您可以参考 Faiss 的 GitHub 问题 获取帮助。

• fairseq：运行 pip install --editable ./fairseq，即可安装修改后的 fairseq 及其依赖包。为保证稳定性，强烈建议使用 Python 3.8。

• 其他依赖：运行 pip install -r requirements.txt。

项目结构

• 预训练 LLM 类（L24，E1024，Alibi 位置编码）：fairseq/fairseq/models/newgpt.py

• 带侧网络的 Transformer 解码器（L12，E1024）：fairseq/fairseq/models/sidenet/transformer_decoder_sidenet.py

• 带侧网络的 Transformer 语言模型类：fairseq/fairseq/models/transformer_lm_sidenet.py

• 记忆库与检索模块：fairseq/fairseq/modules/dynamic_memory_with_chunk.py

• 用于记忆融合的联合注意力机制：fairseq/fairseq/modules/joint_multihead_attention_sum.py

记忆增强的适配训练

数据收集与预处理

请从 官方发布页面 下载 Pile 数据集。Pile 中的每个子数据集都以多个 JSON Lines 文件的形式组织。您可以参考 preprocess/filter_shard_tnlg.py，了解我们如何采样训练集并按照标准的 fairseq 预处理流程进行二进制化处理。

记忆增强的适配训练：

bash train_scripts/train_longmem.sh

评估

请先将预训练的 GPT2-medium 模型和 LongMem 模型的检查点 下载到 checkpoints/ 目录下。

记忆增强的上下文学习

# 评估 GPT2 基线
python eval_scripts/eval_longmem_icl.py --path /path/to/gpt2_pretrained_model
# 评估 LongMem 模型
python eval_scripts/eval_longmem_icl.py --path /path/to/longmem_model --pretrained-model-path /path/to/gpt2_pretrained_model

致谢

LongMem 是基于 fairseq 开发的。同时感谢 eleuther.ai 团队构建了最大规模的高质量语料库——Pile。

LongMem 快速上手指南

LongMem 是一个通过引入长期记忆机制来增强语言模型能力的开源项目。本指南将帮助中国开发者快速完成环境配置、安装及基础评估运行。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统: Linux (推荐 Ubuntu)
• Python 版本: 强烈推荐使用 Python 3.8 以保证稳定性。
• GPU 驱动: 需安装与显卡匹配的 CUDA 驱动。
  • Nvidia V100: 标准 CUDA 环境即可。
  • Nvidia A100/A6000: 需注意 Faiss 兼容性，建议指定 cudatoolkit=11.0。
• PyTorch: 版本需 >= 1.8.0，且必须为 GPU 版本。

安装步骤

请按顺序执行以下命令完成依赖安装和项目配置。为提高下载速度，国内用户可配置 pip 使用清华或阿里镜像源（例如添加 -i https://pypi.tuna.tsinghua.edu.cn/simple）。

1. 安装 PyTorch

请访问 PyTorch 官网 根据您的 CUDA 版本获取安装命令。示例如下：

pip install torch>=1.8.0 torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu111

2. 安装 Faiss-GPU

根据显卡型号选择安装方式：

• V100 显卡:

  pip install faiss-gpu
  

• A100 / A6000 显卡 (需指定 toolkit 版本以避免兼容性问题):

  conda install faiss-gpu cudatoolkit=11.0 -c pytorch
  

3. 安装 Fairseq 及其他依赖

克隆或进入项目目录后，执行以下命令安装修改版的 fairseq 及所需包：

# 安装修改版 fairseq
pip install --editable ./fairseq

# 安装其他依赖包
pip install -r requirements.txt

(注：国内用户可在上述命令后追加 -i https://pypi.tuna.tsinghua.edu.cn/simple)

基本使用

本项目最核心的功能是**记忆增强型上下文学习（Memory-Augmented In-Context Learning）**的评估。以下是基于预训练模型的快速评估示例。

1. 下载预训练模型

首先从 Hugging Face 下载 GPT2-medium 基座模型和 LongMem 模型权重，并存放至 checkpoints/ 目录。

• 模型地址：weizhiwang/LongMem-558M
• 国内加速：建议使用镜像站下载或通过脚本代理下载。

2. 运行评估脚本

评估 GPT2 基线模型：

python eval_scripts/eval_longmem_icl.py --path /path/to/gpt2_pretrained_model

评估 LongMem 模型： 需要同时指定 LongMem 模型路径和对应的 GPT2 预训练模型路径。

python eval_scripts/eval_longmem_icl.py --path /path/to/longmem_model --pretrained-model-path /path/to/gpt2_pretrained_model

提示: 请将 /path/to/... 替换为您本地实际的模型文件路径。