SpikeGPT：基于脉冲神经网络的生成式预训练语言模型

SpikeGPT 是一个轻量级的生成式语言模型，采用了纯二进制、事件驱动的脉冲激活单元。SpikeGPT 的 arXiv 论文可以在这里找到 链接。

如果你对 SpikeGPT 感兴趣，欢迎通过此 链接 加入我们的 Discord 社区！

本仓库的灵感来源于 RWKV-LM。

如果你在环境配置上遇到困难，可以考虑使用我们在 Github 上提供的 SpikeGPT Docker 镜像。

在 Enwik8 数据集上进行训练

1. 请访问以下链接下载 enwik8 数据集： enwik8 数据集。

2. 修改 train.py 脚本中的训练集、验证集和测试集路径，使其指向你解压文件的目录。例如，如果你将文件解压到名为 enwik8_data 的目录下，那么你的 train.py 脚本应更新为如下内容：

   # 设置数据集路径
   datafile_train = "path/to/enwik8_data/train"
   datafile_valid = "path/to/enwik8_data/validate"
   datafile_test = "path/to/enwik8_data/test"
   

在大规模语料库上进行预训练

1. 大规模语料库上的预训练：

   • 首先，对你的语料数据进行预分词处理。
   • 对于自定义数据，可以使用 jsonl2binidx 工具 进行转换。
   • 如果你更倾向于使用已预分词的数据，可以考虑使用预分词的 The Pile，它配备了 20B 参数的分词器，并被 GPT-NeoX 和 Pythia 等模型所采用。
   • 如果资源有限，也可以只使用该数据集中的一份文件，而非整个数据集。

2. 配置训练脚本：

   • 在 train.py 中，取消注释第 82 行，以启用 MMapIndexedDataset 作为数据集类。
   • 将 datafile_train 改为你 binidx 文件的文件名。
   • 注意：不要包含 .bin 或 .idx 文件扩展名。

3. 启动多 GPU 训练：

   • 使用 Hugging Face 的 Accelerate 库，即可开始在多块 GPU 上进行训练。

在 WikiText-103 数据集上进行微调

1. 下载预分词的 WikiText-103 数据集：

   • 你可以从这个 Hugging Face 数据集链接 获取预分词的 WikiText-103 数据集 binidx 文件。

2. 微调模型：

   • 微调模型的方法与预训练时相同。
   • 重要提示：为了避免灾难性遗忘，微调时的学习率应设置得比预训练阶段更低。推荐学习率为 3e-6 左右。
   • 关于批量大小，建议根据具体需求调整，以找到最适合你情况的设置。

基于提示的推理

你可以选择使用自己定制的模型，或者使用我们预训练好的模型来进行推理。我们的预训练模型可以在 这里 获取。该模型曾在 OpenWebText 数据集上进行了 50 亿 token 的训练。

1. 下载我们的预训练模型，并将其放置在本仓库的根目录下。
2. 修改 run.py 中的 'context' 变量为你自定义的提示。
3. 运行 run.py。

结合 NLU 任务进行微调

1. 运行 NLU 文件夹中的相应文件。
2. 将第 17 行中的路径修改为你模型的路径。

引用

如果你在工作中使用了 SpikeGPT，请引用以下文献：

@article{zhu2023spikegpt,
        title = {SpikeGPT: Generative Pre-trained Language Model with Spiking Neural Networks},
        author = {Zhu, Rui-Jie and Zhao, Qihang and Li, Guoqi and Eshraghian, Jason K.},
        journal = {arXiv preprint arXiv:2302.13939},
        year    = {2023}
}

SpikeGPT 快速上手指南

SpikeGPT 是一个轻量级的生成式语言模型，采用纯二进制、事件驱动的脉冲神经网络（SNN）激活单元。本指南将帮助你快速完成环境配置、训练及推理。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu)
• 硬件: 支持 CUDA 的 NVIDIA GPU（用于多卡训练加速）
• Python: 3.8+

前置依赖

项目依赖 PyTorch 及相关深度学习库。如果你在安装过程中遇到环境配置困难，强烈推荐使用社区提供的 Docker 镜像：

• Docker 镜像: SpikeGPT-container

若手动安装，请确保已安装 PyTorch（带 CUDA 支持）。

安装步骤

1. 克隆仓库

   git clone https://github.com/ridgerchu/SpikeGPT.git
   cd SpikeGPT
   

2. 安装 Python 依赖

   pip install -r requirements.txt
   

   注：若 requirements.txt 不存在，请根据项目代码安装 torch, accelerate, transformers 等核心库。

3. （可选）使用 Docker 如果本地环境配置复杂，可直接拉取容器运行：

   docker pull eddiem3/spikegpt-container
   docker run --gpus all -it eddiem3/spikegpt-container
   

基本使用

1. 快速推理 (Inference)

使用官方预训练模型进行文本生成是最简单的上手方式。该模型在 OpenWebText 数据集上训练了 50 亿 token。

步骤：

1. 下载预训练模型 从 Hugging Face 下载模型文件并放入项目根目录： SpikeGPT-OpenWebText-216M

2. 修改提示词 编辑 run.py 文件，修改 context 变量为你想要的输入提示：

   # 在 run.py 中
   context = "Once upon a time,"  # 修改此处
   

3. 运行推理

   python run.py
   

2. 小规模数据训练 (Enwik8)

适合快速验证模型训练流程。

1. 下载数据集 访问 enwik8 dataset 下载数据。

2. 配置路径 解压数据后，修改 train.py 中的数据集路径：

   # 在 train.py 中
   datafile_train = "path/to/enwik8_data/train"
   datafile_valid = "path/to/enwik8_data/validate"
   datafile_test = "path/to/enwik8_data/test"
   

3. 开始训练

   python train.py
   

3. 大规模预训练与微调

预训练 (Pre-training)

1. 数据准备: 使用 jsonl2binidx tool 将自定义数据转换为 binidx 格式，或直接使用预处理的 The Pile 数据集。
2. 脚本配置: 在 train.py 中取消第 82 行注释以启用 MMapIndexedDataset，并将 datafile_train 指向你的 binidx 文件名（不要包含 .bin 或 .idx 后缀）。
3. 多卡启动:

   accelerate launch train.py
   

微调 (Fine-Tuning on WikiText-103)

1. 下载数据: 获取预处理好的 WikiText-103 binidx。
2. 调整参数:
   • 复用预训练脚本逻辑。
   • 关键: 设置较小的学习率（推荐 3e-6）以避免灾难性遗忘。
   • 根据显存调整 batch size。

NLU 任务微调

进入 NLU 文件夹，修改对应脚本第 17 行的模型路径后直接运行即可。