TokenFormer：一个完全基于注意力机制的神经网络，具有令牌化的模型参数。通过令牌化一切来最大化 Transformer 的灵活性。

本仓库是我们论文 TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters 及其后续工作的官方实现。我们的 TokenFormer 是一种原生可扩展的架构，它不仅利用注意力机制进行输入令牌（tokens）之间的计算，还用于令牌与模型参数之间的交互，从而增强了架构的灵活性。我们已尽最大努力确保代码库干净、简洁、易于阅读、保持最先进水平，并且仅依赖于最少的依赖项。

TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters

Haiyang Wang*, Yue Fan*, Muhammad Ferjad Naeem, Yongqin Xian, Jan Eric Lenssen, Liwei Wang, Federico Tombari, Bernt Schiele

• 主要联系人：Haiyang Wang (haiwang@mpi-inf.mpg.de)(wanghaiyang6@stu.pku.edu.cn), Bernt Schiele (schiele@mpi-inf.mpg.de)

📣 新闻

• [25-02-11] 🔥 TokenFormer 被接受为 亮点报告。
• [25-01-22] 🔥 TokenFormer 被 ICLR2025 接收。
• [25-01-12] 在 TPU (GCP-Cloud) 上的 Jax 代码已发布，请参见 此处。
• [24-11-08] 🚀 使用 PyTorch 的训练代码已发布。
• [24-11-02] 如果我遗漏了任何相关论文，请随时通过电子邮件联系我。我将尽力在未来版本中包含所有相关论文。
• [24-10-31] 🚀 使用 PyTorch 的推理代码已发布。
• [24-10-31] 👀 TokenFormer 在 arXiv 上发布。

🔥 一些思考

• 我们旨在为模型提供一个新视角，适用于未来的任何计算图。理论上，通过使用数据令牌、参数令牌和记忆令牌，并通过点积交互，可以灵活地构建任何网络。这里有很多设计可能性。例如，引入记忆令牌可以构建类似于 Mamba 的类 RNN 网络。将参数令牌与记忆令牌合并可以创建类似于 TTT 网络 的东西。参数令牌也可以反向关注输入数据，使网络参数动态地依赖于数据，逐层更新。

概述

• 💫 我们的目标
• 🚀 主要结果
• 📘 模型库
• 🛠️ 快速开始
• 👀 待办事项
• 📘 引用

💫 我们的目标？

我们介绍了 Tokenformer，一种 完全基于注意力机制 的架构，它通过完全采用注意力机制统一了令牌-令牌和令牌-参数交互的计算，最大化神经网络的灵活性。这一优势使其能够处理可变数量的参数，本质上增强了模型的可扩展性，促进了渐进式高效扩展。

我们不仅对数据进行令牌化，还对模型参数进行令牌化，用数据和参数令牌之间的交互流取代了模型的概念，进一步推动了网络架构的统一。

希望这种架构能够比传统 Transformer 提供更大的灵活性，并将进一步促进基础模型、稀疏推理（MoE）、参数高效调优、设备-云协作、视觉-语言、模型可解释性等领域的发展。

# 给定输入的 Pattention 实现

query, key, value = inputs, key_param_tokens, value_param_tokens

attn_weight = query @ key.transpose(-2, -1) * scale_factor

attn_weight *= attn_masks
# 修改后的 softmax，softmax 等同于 exp + L1 范数
attn_weight = nonlinear_norm_func(attn_weight, self.norm_activation_type, dim=-1)

output = attn_weight @ value

🚀 主要结果

增量模型扩展

传统上，大型 Transformer 架构是从头开始训练的，没有重用先前较小规模的模型。在本文中，我们提出了一种新颖的完全基于注意力机制的架构，允许增量扩展模型，从而大大降低了训练大型 Transformer 架构的总体成本。

在 Pile 数据集上进行语言建模并进行零样本评估

(零样本评估。) 每个模型大小的最佳性能以粗体突出显示。我们的比较是与使用各种分词器的公开可用的基于 Transformer 的语言模型进行的。遵循 Pythia，我们的模型在 Pile 数据集上训练了高达 300B 个令牌。

在 ImageNet-1k 分类上进行视觉建模

(图像分类。) 在 ImageNet-1K 上标准视觉 Transformer 的比较。

📘 模型库

语言建模基准测试（Pile）

预训练模型已上传至 huggingface，包括 TokenFormer-150M、TokenFormer-450M、TokenFormer-900M 和 TokenFormer-1-5B，这些模型在 Pile 数据集上训练了 300B 个 token。

这些模型在 Pile 数据集上训练，遵循 Transformer 的标准模型维度，并在 mamba 描述的零样本（zero-shot）基准测试上进行了评估：

模型	参数	层数	模型维度	检查点	配置	日志
TokenFormer-150M	150M	12	768	检查点	配置	日志
TokenFormer-450M	450M	24	1024	检查点	配置	日志
TokenFormer-900M	900M	32	1280	检查点	配置	日志
TokenFormer-1-5B	1-5B	40	1536	检查点	配置	日志

注意：这些是仅训练了 300B token 的基础模型，未经过任何形式的下游修改（如指令微调等）。预期性能与其他在相似数据上训练的架构相当或更优，但无法与更大规模或经过微调的模型相比。

视觉建模基准测试（基于 CLIP 方法的 DataComp-1B）

稍后发布。

🛠️ 快速开始

安装

首先确保您处于 Python 3.8 和 CUDA 12 的环境中，并安装了适当版本的 PyTorch 1.8 或更高版本。注意： 我们的 TokenFormer 基于 GPT-NeoX，GPT-NeoX 依赖的一些库尚未更新以兼容 Python 3.10+。Python 3.9 似乎可以工作，但本代码库是在 Python 3.8 上开发和测试的。

要安装其余的基本依赖项，请运行：

conda create -n TokenFormer python=3.8

git clone https://github.com/Haiyang-W/TokenFormer.git

pip install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 --index-url https://download.pytorch.org/whl/cu121

### raven module load gcc/10

### 如果在运行 pip install -r requirements/requirements.txt 时遇到 cargo 问题，请执行以下命令
# curl https://sh.rustup.rs -sSf | sh
# export PATH="$HOME/.cargo/bin:$PATH"
# source ~/.profile
# source ~/.cargo/env

### 如果在运行 pip install -r requirements/requirements.txt 时遇到 mpi4py 问题，请执行：
# conda install -c conda-forge mpi4py=3.0.3

pip install -r requirements/requirements.txt

pip install -r requirements/requirements-flashattention.txt # 需要 gcc > 9
pip install -r requirements/requirements-wandb.txt # 可选，如果使用 WandB 记录日志
pip install -r requirements/requirements-tensorboard.txt # 可选，如果使用 tensorboard 记录日志
pip install -r requirements/requirements-comet.txt # 可选，如果使用 Comet 记录日志

# 安装 apex
pip install -r requirements/requirements-apex-pip.txt # pip > 23.1
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./

从仓库根目录执行。

评估

要运行模型的零样本评估（对应论文中的表 1），我们使用 lm-evaluation-harness 库。

首先，您需要从 huggingface 下载预训练权重到本地目录。例如，相对于仓库根目录的路径是 ./TokenFormer-150M/pytorch_model.bin。

# 单 GPU 评估（目前仅在单 GPU 上测试过。）

cd ./TokenFormer
python ./deepy.py eval.py -d configs tokenformer/150M_eval.yml --eval_tasks lambada_openai hellaswag piqa arc_challenge arc_easy winogrande

用于训练的预配置数据集

提供了几个预配置的数据集，包括来自 openwebtext 和 Pile 的大部分组件。

例如，使用 GPT-NeoX 20B Tokenizer 下载并分词 openwebtext2 数据集。您可以先尝试这个小数据集。

python prepare_data.py -d ./data -t HFTokenizer --vocab-file tokenizer.json openwebtext2

预处理后的数据将位于 ./data/openwebtext2。

对于 Pile 300B（非官方且未复制的版本）：

python prepare_data.py -d ./data -t HFTokenizer --vocab-file tokenizer.json pile

预处理后的数据将位于 ./data/pile。

分词后的数据将保存为两个文件：[data-dir]/[dataset-name]/[dataset-name]_text_document.bin 和 [data-dir]/[dataset-name]/[dataset-name]_text_document.idx。您需要在训练配置文件的 data-path 字段中添加这两个文件共享的前缀。例如：

"data-path": "./data/pile/pile_0.87_deduped_text_document",

如果您想轻松上手，可以尝试 enwik8。

训练

单节点启动

注意，这是针对单节点的。适用于您可以直接 SSH 登录到一台 8-GPU 机器并直接运行程序的情况。

CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python deepy.py train.py configs/tokenformer/150M_train_pile.yml

多节点启动

请参考 multi-node-launching。我使用 slurm 并提供以下指导。

首先，修改您的训练配置

{
  "launcher": "slurm",
  "deepspeed_slurm": true,
}

然后，我提供一个使用 16 个 GPU 的 slurm 脚本作为示例。

#!/bin/bash
#SBATCH --job-name="150M_16gpus"
#SBATCH --constraint="gpu"
#SBATCH --nodes=2
#SBATCH --ntasks-per-node=8
#SBATCH --gres=gpu:8
#SBATCH --cpus-per-task=4   #   每个任务使用 4 个核心。
#SBATCH --time=24:00:00
#SBATCH -o /tmp/150M_%A_%a.out

conda activate TokenFormer

一些可能用到的分布式环境变量

export HOSTNAMES=scontrol show hostnames "$SLURM_JOB_NODELIST" export MASTER_ADDR=$(scontrol show hostnames "$SLURM_JOB_NODELIST" | head -n 1) export MASTER_PORT=12856 export COUNT_NODE=scontrol show hostnames "$SLURM_JOB_NODELIST" | wc -l

运行之前编写的 hostfile 创建脚本

bash ./write_hostfile.sh

通过 DLTS_HOSTFILE 告诉 DeepSpeed 在哪里找到我们生成的 hostfile，你可以自定义任何路径。

export DLTS_HOSTFILE=/tmp/hosts_$SLURM_JOBID

python3 deepy.py train.py ./configs/tokenformer/150M_train_pile.yml

这里的所有路径都可以自定义；你可以将上述脚本和 `write_hostfile.sh` 中的 `/tmp` 替换为你想要的任何路径。然后运行脚本

sbatch scripts.sh

#### 训练后的零样本评估
进入你的检查点目录，例如 150M

cd ./work_dirs/150M_TokenFormer_Pile/checkpoints python zero_to_fp32.py . pytorch_model.bin

然后用该路径替换 `150M_eval.yml` 中的 **[eval_ckpt](https://github.com/Haiyang-W/TokenFormer/blob/79a02d8a2f847e8bbc627f7cb1632a2f24f3f826/configs/tokenformer/150M_eval.yml#L98)**。

cd ./TokenFormer python ./deepy.py eval.py -d configs tokenformer/150M_eval.yml --eval_tasks lambada_openai hellaswag piqa arc_challenge arc_easy winogrande

`注意：` 我目前只运行了前 1000 次迭代的训练代码来检查损失，看起来没问题，所以先发布出来供大家使用。我不能保证完全没有问题。如果你愿意等待，我可以做最终检查，但这可能需要一些时间。

### 增量训练
请从 [huggingface:354M_TokenFormer_Openwebtext2.zip](https://huggingface.co/Haiyang-W/TokenFormer-354M-Openwebtext2/tree/main) 下载在 openwebtext2 数据集上预训练的 354M TokenFormer，并将其解压到你的本地目录。例如，相对于仓库根目录的相对路径是 ``./354M_TokenFormer_Openwebtext2/``。然后，将 [语言数据集](https://github.com/Haiyang-W/TokenFormer/blob/ab93c3de1805a6a7a2733c462068d68e83c849b5/configs/incremental_scaling_openwebtext2/354M_to_757M_train_openwebtext2_60k.yml#L110) 和 [预训练检查点](https://github.com/Haiyang-W/TokenFormer/blob/ab93c3de1805a6a7a2733c462068d68e83c849b5/configs/incremental_scaling_openwebtext2/354M_to_757M_train_openwebtext2_60k.yml#L33) 的路径正确配置为你的本地路径。这里的预训练检查点是在 OpenWebText2 上训练 TokenFormer 354M 模型 60 万步的结果。对应的配置文件在[这里](https://github.com/Haiyang-W/TokenFormer/blob/main/configs/incremental_scaling_openwebtext2/354M_train_openwebtext2_basemodel.yml)。

请按照[此处](https://github.com/Haiyang-W/TokenFormer/tree/main?tab=readme-ov-file#preconfigured-datasets-for-training)的说明准备 openwebtext2 数据集。

cd ./TokenFormer python deepy.py train.py configs/incremental_scaling_openwebtext2/354M_to_757M_train_openwebtext2_60k.yml

模型的性能如下：
|  模型  | 训练策略 | 层数 | 模型维度 | 迭代次数 |验证集困惑度|配置文件|
|---------|---------|---------|--------|--------|---------|---------|
|  TokenFormer-354M | 从头训练 | 24 | 1024  | 600k |11.9|[配置文件](https://github.com/Haiyang-W/TokenFormer/blob/main/configs/incremental_scaling_openwebtext2/354M_train_openwebtext2_basemodel.yml) |
|  TransFormer-757M | 从头训练 | 24 | 1536  | 60k |12.0 | - |
|  TransFormer-757M | 从头训练 | 24 | 1536  | 120k |11.3 | - |
|  TransFormer-757M | 从头训练 | 24 | 1536  | 600k |10.5 | - |
|  TokenFormer-757M | 增量训练 | 24 | 1024 | 60k |10.9|[配置文件](https://github.com/Haiyang-W/TokenFormer/blob/main/configs/incremental_scaling_openwebtext2/354M_to_757M_train_openwebtext2_60k.yml) |
|  TokenFormer-757M | 增量训练 | 24 | 1024 | 120k |10.7|[配置文件](https://github.com/Haiyang-W/TokenFormer/blob/main/configs/incremental_scaling_openwebtext2/354M_to_757M_train_openwebtext2_120k.yml) |

## 👀 待办事项

- [x] 发布 [arXiv](https://arxiv.org/abs/2410.23168) 版本。
- [x] 发布大语言模型的推理代码和模型权重。
- [x] 发布大语言模型的训练代码。
- [x] 发布大语言模型的增量缩放训练代码。

## 📘 引用
如果我们的工作对你有帮助，请考虑按如下方式引用。

@inproceedings{ wang2025tokenformer, title={TokenFormer: Rethinking Transformer Scaling with Tokenized Model Parameters}, author={Haiyang Wang and Yue Fan and Muhammad Ferjad Naeem and Liwei Wang and Yongqin Xian and Jan Eric Lenssen and Federico Tombari and Bernt Schiele}, booktitle={The Thirteenth International Conference on Learning Representations}, year={2025}, url={https://openreview.net/forum?id=oQ4igHyh3N} } ```

TokenFormer 快速上手指南

TokenFormer 是一个完全基于注意力机制的神经网络，通过将模型参数也“令牌化”，最大化 Transformer 架构的灵活性。本指南将帮助您快速搭建环境并运行 TokenFormer。

环境准备

系统要求：

• Python 版本： 3.8（首选）。代码主要基于 Python 3.8 开发和测试，Python 3.9 可能可用，但不保证兼容性。Python 3.10+ 因依赖库限制暂不支持。
• CUDA 版本： 12.x
• PyTorch 版本： 1.8 或更高版本（推荐 2.2.1）
• 其他： 需要安装 gcc（版本 > 9）以支持 FlashAttention。

安装步骤

1. 创建并激活 Conda 环境：

   conda create -n TokenFormer python=3.8
   conda activate TokenFormer
   

2. 克隆仓库并安装 PyTorch：

   git clone https://github.com/Haiyang-W/TokenFormer.git
   cd TokenFormer
   pip install torch==2.2.1 torchvision==0.17.1 torchaudio==2.2.1 --index-url https://download.pytorch.org/whl/cu121
   

3. 安装基础依赖：

   pip install -r requirements/requirements.txt
   

4. 安装 FlashAttention 等可选依赖：

   # 安装 FlashAttention (需要 gcc > 9)
   pip install -r requirements/requirements-flashattention.txt
   # 以下为可选日志工具，按需安装
   pip install -r requirements/requirements-wandb.txt
   pip install -r requirements/requirements-tensorboard.txt
   pip install -r requirements/requirements-comet.txt
   

5. 安装 NVIDIA Apex（用于混合精度训练）：

   pip install -r requirements/requirements-apex-pip.txt
   git clone https://github.com/NVIDIA/apex
   cd apex
   pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings "--build-option=--cpp_ext" --config-settings "--build-option=--cuda_ext" ./
   cd ..
   

安装问题提示：

• 若安装依赖时遇到 cargo 相关问题，请先安装 Rust：

  curl https://sh.rustup.rs -sSf | sh
  export PATH="$HOME/.cargo/bin:$PATH"
  source ~/.profile
  source ~/.cargo/env
  

• 若遇到 mpi4py 安装问题，可尝试通过 Conda 安装：

  conda install -c conda-forge mpi4py=3.0.3
  

基本使用

1. 模型推理与评估

首先，从 Hugging Face 下载预训练模型权重（例如 TokenFormer-150M）到本地目录，例如 ./TokenFormer-150M/pytorch_model.bin。

然后，使用以下命令进行零样本评估（单 GPU）：

cd ./TokenFormer
python ./deepy.py eval.py -d configs tokenformer/150M_eval.yml --eval_tasks lambada_openai hellaswag piqa arc_challenge arc_easy winogrande

2. 数据准备（用于训练）

TokenFormer 支持多个数据集。以下以准备 openwebtext2 数据集为例：

python prepare_data.py -d ./data -t HFTokenizer --vocab-file tokenizer.json openwebtext2

预处理后的数据将保存在 ./data/openwebtext2 目录下。

若要准备用于 300B 令牌训练的 Pile 数据集（非官方无版权版本）：

python prepare_data.py -d ./data -t HFTokenizer --vocab-file tokenizer.json pile

预处理后的数据将保存在 ./data/pile 目录下。

3. 模型训练

准备好数据后，可以使用以下命令启动训练。以 TokenFormer-150M 模型配置为例：

python ./deepy.py train.py -d configs tokenformer/150M.yml

注意： 训练脚本支持多种配置，您可以根据需要修改对应的 YAML 配置文件（位于 configs/tokenformer/ 目录下）来调整模型规模、超参数和数据路径。