Samba：用于高效无限上下文语言建模的简单混合状态空间模型

Samba 是一种简单而强大的混合模型，具有无限的上下文长度。它的架构极其简洁：

Samba = Mamba + MLP + 滑动窗口注意力 + 层级 MLP 堆叠。

我们最大的模型 Samba-3.8B 在 Phi3 数据集的 3.2 万亿个 token 上进行训练，在多个主流基准测试中（如 MMLU、GSM8K 和 HumanEval）大幅超越了 Phi3-mini。此外，Samba 只需少量指令微调即可实现完美的长上下文检索能力，同时保持与序列长度呈线性关系的计算复杂度。这一特性使得 Samba-3.8B-instruct 在长上下文摘要等下游任务中表现出色。

性能 :rocket:

模型	MMLU	GSM8K	HumanEval	GovReport	SQuALITY
Phi-3-mini-4K-instruct	68.8	82.5	58.5	14.4	21.6
Samba-3.8B-instruct (预览)	71.9	87.6	62.8	18.9	21.2

我们在 MMLU 报告 5-shot 准确率，GSM8K 报告 8-shot CoT 准确率，HumanEval 报告 0-shot pass@1，GovReport 和 SQuALITY 则报告 ROUGE-L 分数。

更新

• [7月9日] 我们发布了基于 SambaY 架构的 Phi-4-mini-flash-reasoning 模型。请查看新的代码库中的神经规模实验 这里!
• [1月22日] Samba 已被 ICLR 2025 接受！
• [12月8日] 添加了评估脚本和更多基线架构。
• [6月11日] 发布了在 SlimPajama 数据集上训练 Samba-421M 和 Samba-1.3B 的代码库。

代码概览

我们在 SlimPajama 数据集上的训练基础设施是对 TinyLlama 和 LitGPT 的修改版本。用户可以通过修改 model_name 和包含论文中提到的大量基线架构的 config 文件，轻松指定不同的架构配置。我们的 RetNet 和 GLA 实现来自优秀的 Flash Linear Attention 仓库。

从头开始预训练 Samba

请按照 Dockerfile 设置环境。数据准备主要遵循 TinyLlama，但我们仅使用 SlimPajama 数据集。

数据准备

将 Slimpajama 数据集下载到您选择的目录。我们还提供了预分词的数据 这里。

cd /path/to/dataset
git lfs install
git clone https://huggingface.co/datasets/cerebras/SlimPajama-627B

SlimPajama 数据集需要 893GB 的磁盘空间。使用提供的脚本对数据集进行分词并分成若干块。

python scripts/prepare_slimpajama.py --source_path /path/to/SlimPajama --tokenizer_path data/llama  --destination_path data/slim --split validation --percentage 1.0
python scripts/prepare_slimpajama.py --source_path /path/to/SlimPajama --tokenizer_path data/llama  --destination_path data/slim --split train --percentage 1.0

现在您已经准备好启动训练任务了！

训练

以下脚本在一个拥有 8 张 GPU 的单节点上训练一个默认的 Samba-421M 模型，使用 200 亿个 token。

torchrun --nnodes=1 --nproc_per_node=8 --rdzv_id=samba-421M --rdzv_backend=c10d  --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} pretrain.py --train_data_dir data/slim --val_data_dir data/slim 

您可以将 model_name 修改为 "Samba_1.3B"，并将 train_config 修改为 "tsz512x4k_100B"，以训练一个拥有 1000 亿个 token 的 Samba-1.3B 模型。我们假设您有 8 个节点，每个节点配备 8 张 GPU；您也可以修改 nodes 的数量来减少使用的 GPU 数量。

评估

我们利用 lm-evaluation-harness 对预训练模型进行评估。目前我们仅支持非生成类任务。

pip install lm-eval
python eval.py --model Samba \
          --model_args pretrained=path/to/ckpt.pth,config="Samba_1.3B" \
          --tasks lambada_openai,arc_easy,winogrande,hellaswag,piqa --device cuda:0 --batch_size 1 --trust_remote_code 

引用

如果您觉得我们的工作有用，请考虑引用：

@article{ren2024samba,
      title={Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling}, 
      author={Liliang Ren and Yang Liu and Yadong Lu and Yelong Shen and Chen Liang and Weizhu Chen},
      journal = {arXiv preprint},
      year={2024},
      url={https://arxiv.org/abs/2406.07522}
}

联系方式

Liliang Ren (liliangren@microsoft.com)

Samba 快速上手指南

Samba 是一种简单而强大的混合架构模型，结合了 Mamba、MLP 和滑动窗口注意力机制，支持无限上下文长度且保持线性复杂度。本指南将帮助你从零开始准备环境并训练 Samba 模型。

环境准备

系统要求

• 操作系统: Linux (推荐 Ubuntu)
• GPU: 支持 CUDA 的 NVIDIA GPU (训练示例需多卡环境，如 8x A100/V100)
• 磁盘空间: 至少 1TB (SlimPajama 数据集约占用 893GB)
• Docker: 推荐使用官方提供的 Dockerfile 构建环境

前置依赖

项目基于 TinyLlama 和 LitGPT 修改而成。主要依赖包括：

• PyTorch
• Flash Linear Attention (用于 RetNet 和 GLA 实现)
• lm-evaluation-harness (用于评估)

建议直接通过项目根目录的 Dockerfile 搭建环境以确保依赖一致性：

docker build -t samba-env .

安装步骤

1. 获取代码

克隆仓库并进入目录：

git clone https://github.com/microsoft/Samba.git
cd Samba

2. 数据准备

Samba 训练主要使用 SlimPajama 数据集。你可以下载原始数据或使用已分词的数据。

方案 A：下载原始数据并分词（推荐）

# 安装 Git LFS
git lfs install

# 下载数据集 (约 893GB)
cd /path/to/dataset
git clone https://huggingface.co/datasets/cerebras/SlimPajama-627B

# 返回项目根目录进行分词处理
# 注意：请替换 /path/to/SlimPajama 为实际路径
python scripts/prepare_slimpajama.py --source_path /path/to/SlimPajama --tokenizer_path data/llama  --destination_path data/slim --split validation --percentage 1.0
python scripts/prepare_slimpajama.py --source_path /path/to/SlimPajama --tokenizer_path data/llama  --destination_path data/slim --split train --percentage 1.0

方案 B：直接使用预分词数据（加速） 如果网络受限，可直接从 Hugging Face 下载已处理好的数据：

git lfs install
git clone https://huggingface.co/datasets/jsun/slimpajama_Llama2_Tokenizer
# 将下载的数据链接或移动到 project/data/slim 目录

基本使用

1. 启动训练

以下命令将在单节点 8 卡环境下训练默认的 Samba-421M 模型（20B tokens）。

torchrun --nnodes=1 --nproc_per_node=8 --rdzv_id=samba-421M --rdzv_backend=c10d  --rdzv_endpoint=${MASTER_ADDR}:${MASTER_PORT} pretrain.py --train_data_dir data/slim --val_data_dir data/slim 

调整模型大小： 若要训练 Samba-1.3B 模型（100B tokens），请修改 pretrain.py 中的配置：

• 将 model_name (第 33 行) 改为 "Samba_1.3B"
• 将 train_config (第 34 行) 改为 "tsz512x4k_100B"
• 根据显卡数量调整 nodes (第 43 行)，默认假设 8 节点 x 8 卡。

2. 模型评估

使用 lm-evaluation-harness 对预训练模型进行评估（当前仅支持非生成式任务）。

首先安装评估工具：

pip install lm-eval

运行评估脚本：

python eval.py --model Samba \
          --model_args pretrained=path/to/ckpt.pth,config="Samba_1.3B" \
          --tasks lambada_openai,arc_easy,winogrande,hellaswag,piqa --device cuda:0 --batch_size 1 --trust_remote_code 

请将 path/to/ckpt.pth 替换为你训练完成的检查点路径。