RetNet

基于 PyTorch 的 Retentive Network：大型语言模型的 Transformer 替代方案 实现。

关于本仓库

这是一个极简、纯 PyTorch 实现的 RetNet。RetNet 论文：Retentive Network：大型语言模型的 Transformer 替代方案。

本仓库的贡献者并非原论文的作者。RetNet 的思想与公式设计完全归功于原作者。

本仓库的目的是促进科学与技术的理解与进步。代码在优先考虑正确性和可读性的同时，暂未进行高度优化。

已实现的功能

• 单尺度与多尺度保持机制：
  • 并行范式
  • 递归范式
  • 分块范式
• 带有 FFN 和 LayerNorm 的多层保持网络：
  • 并行范式
  • 递归范式
  • 分块范式
• 构建于保持网络之上的因果语言模型（CLM）

使用与示例：

• 请参阅以 test_ 为前缀的脚本，了解基本用法示例。

位置编码

src/ 中的主要实现使用了 Microsoft 的 xPos 进行位置编码。

src/complex 中的实现则采用复数值来编码位置，这要求参数和数据的类型为 torch.ComplexFloat（64 位）。由于目前 PyTorch 尚未支持半精度复数类型，这一实现存在一些限制；此外，它所需的内存是 32 位精度实数值的两倍。

贡献

欢迎所有贡献。请查看 issues，了解当前需要完成的工作。

如果您希望为本项目贡献力量，请先 Fork 本仓库，然后提交 Pull Request 以供审核。

参考文献

@misc{sun2023retentive,
      title={Retentive Network: A Successor to Transformer for Large Language Models}, 
      author={Yutao Sun and Li Dong and Shaohan Huang and Shuming Ma and Yuqing Xia and Jilong Xue and Jianyong Wang and Furu Wei},
      year={2023},
      eprint={2307.08621},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}

RetNet 快速上手指南

RetNet 是 Transformer 架构的潜在继任者，本仓库提供了基于 PyTorch 的极简、纯净实现，优先保证代码的正确性与可读性。

环境准备

• 系统要求：Linux / macOS / Windows
• Python 版本：建议 Python 3.8+
• 核心依赖：
  • PyTorch (最新稳定版)
  • torchscale (用于 xPos 位置编码，可选但推荐)

注意：若使用 src/complex 目录下的复数实现，需确保数据类型为 torch.ComplexFloat (64-bit)，目前 PyTorch 尚不支持半精度复数类型，且内存占用约为实数类型的两倍。

安装步骤

1. 克隆仓库

   git clone https://github.com/Jamie-Stirling/RetNet.git
   cd RetNet
   

2. 安装依赖 推荐使用国内镜像源加速安装：

   pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
   # 如果需要使用 xPos 位置编码，安装 torchscale
   pip install torchscale -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   (注：若项目中包含 requirements.txt，请优先执行 pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple)

基本使用

本仓库未提供封装好的 CLI 工具，主要通过运行测试脚本来验证功能和作为使用示例。

1. 查看示例脚本 进入项目根目录，查找以 test_ 开头的脚本：

   ls scripts/test_*.py
   # 或直接在 src 目录下查找
   find . -name "test_*.py"
   

2. 运行基础示例 运行任意测试脚本即可看到 RetNet 的基本调用方式（包括并行、递归和分块范式）：

   python scripts/test_retention.py
   

   具体文件名请以实际仓库中的 scripts 或根目录下的文件为准。

3. 代码集成示例 在你的 Python 项目中导入并使用核心模块：

   import torch
   from src.retention import MultiScaleRetention # 假设路径，具体请参考 test 脚本导入方式
   
   # 初始化模型参数
   model = MultiScaleRetention(
       dim=512,
       heads=8,
       chunk_size=64
   )
   
   # 准备输入 (batch, seq_len, dim)
   x = torch.randn(2, 1024, 512)
   
   # 前向传播 (平行范式示例)
   output = model(x)
   print(output.shape)
   

提示：由于代码优先追求可读性而非极致优化，建议先阅读 scripts/ 目录下的 test_ 系列文件，它们是理解不同范式（Parallel, Recurrent, Chunkwise）用法的最佳文档。