Pythia: 跨时间与规模解读 Transformer

本仓库为 EleutherAI 的 Pythia 项目代码。该项目将可解释性分析与缩放定律相结合，旨在理解自回归 Transformer 在训练过程中知识的发展与演变。关于模型、训练过程及其特性的详细信息，请参阅我们的论文 Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling。

Pythia 套件的开发明确旨在解决现有模型套件在可解释性、学习动态、伦理与透明度研究方面的不足。其核心特性包括：

1. 完全公开：论文中使用的所有模型、数据和代码均已公开发布，确保结果可完全复现。论文中的所有结果均至少已由另一个实验室独立验证。
2. 丰富的训练检查点：所有模型在训练过程中均保存了 154 个检查点，便于研究 LLM 的学习动态。
3. 数据顺序一致：所有模型均使用相同的数据、按相同的顺序进行训练，使研究人员能够探索对训练过程的因果干预。

在发布时，Pythia 是世界上唯一满足这些需求的模型套件。事实上，我们为 12B 参数模型发布的 154 个检查点，其数量超过了当时全球所有其他 12B+ 参数模型组合所发布的部分训练检查点总数。我们的工作已启发其他团队创建类似项目，包括 LLM360 的 Amber 和 K2-65B、AI2 的 OLMo 以及 Zyphra 的 BlackMamba。

除 Pythia 套件本身外，本仓库还作为信息枢纽，包含以下论文的相关信息、代码和复现说明：

• 大型语言模型中的涌现与可预测记忆 [代码]
• PolyPythias: 五十次语言模型预训练运行中的稳定性与异常值 [代码]

更新日志

• [2025年3月10日] 添加了 PolyPythias 论文的信息。
• [2024年7月9日] 大幅修订了 README，包括更好的历史背景说明，并推广了许多基于 Pythia 的优秀研究工作。同时添加了后续训练模型的链接。
• [2023年11月2日] 应部分研究者请求，新增了 14M 和 31M 参数模型。我们计划在未来训练这些模型的去重版本。
• [2023年4月3日] 发布了所有 Pythia 模型的新版本，修复了原始套件中的各种不一致问题。有关更改的详细信息，请参阅 Pythia 论文 的附录 B。旧模型（"v0"）仍保留在此处，可用于消融研究。
• [2023年1月20日] 我们将 Pythia 模型套件重命名，使其总数参数包含嵌入层和去嵌入层参数，这与许多其他模型套件保持一致，我们相信此惯例更能反映这些模型在设备上的内存使用情况。我们还发现，由于一个拼写错误，其中一个模型比预期要小，因此将其替换为预期大小的模型。详见此处。

目录

• 模型
  • 多种随机种子
• 使用 Pythia
  • 快速开始
  • 复现训练
  • 探索数据集
  • Pythia 论文复现
• 基准测试分数
• 基于 Pythia 的研究
  • 语言模型内部机制
  • 学习动态
  • 训练数据如何决定模型行为
  • 安全、审计与合规性研究
• 引用详情
• 许可证

模型

我们在 Pile 数据集（论文, 数据表）以及去重后的 Pile 数据集上训练并发布了一系列共 8 个不同规模的模型。所有 8 个模型均使用完全相同的数据、以完全相同的顺序进行训练。每个模型在训练期间观察了约 3000 亿个 token。对于“标准”模型，这对应 Pile 上略低于 1 个 epoch；对于去重后的 Pile（1 个 epoch 包含 2070 亿个 token），则对应约 1.5 个 epoch。所有模型均使用混合精度训练，除 EleutherAI/pythia-1b 因在 fp16 下后期出现无法解决的损失尖峰而使用 bf16 外，其余均使用 fp16。

在初始发布之后，应从事稀疏自编码器缩放研究的对齐研究人员的要求，我们训练了 14M 和 31M 参数的模型。

为了促进 LLM 学习动态的研究，我们为每个模型提供了 154 个检查点，对应于步骤 0（初始化），1，2，4，8，16，32，64，128，256，512，1000，以及之后每 1000 步。我们还上传了预分词的数据文件和一个用于重建训练期间数据加载器的脚本。详见复现训练部分。

用于使用 GPT-NeoX 库 训练这些模型的配置文件可以在此仓库的 models/ 目录以及 GPT-NeoX 库本身中找到。

我们在最初训练这些模型时犯了一个错误，导致不同运行之间存在一些不一致。我们修复了这些不一致并重新运行了整个模型套件，原始运行可在名称 EleutherAI/pythia-160m-v0 下找到。有关 v0 模型与主套件差异的更多详细信息，请参阅 Pythia 论文。

所有模型的损失曲线包含在我们（杂乱的！）wandb 项目中，链接为这里。

模型与 wandb 运行之间大致的部分对应关系如下：

多种随机种子

用于训练 Pythia 模型的随机种子是 GPT-NeoX 的默认值：1234。为了研究随机性如何影响模型行为，我们一直在使用不同的随机种子训练更多模型。目前，我们已经训练并发布了以下使用种子 1 到 9 的模型。

• Pythia 14M
• Pythia 31M
• Pythia 70M
• Pythia 160M
• Pythia 410M

所有这些都是 标准 Pythia 模型，而不是在去重 Pile 上训练的模型。与最初发布的模型结合，它们代表了十个使用不同随机种子、其他方面完全相同的变体。可以在 HuggingFace 上使用命名模式 https://huggingface.co/EleutherAI/pythia-[size]-seed[num] 找到它们。例如，https://huggingface.co/EleutherAI/pythia-160m-seed7。请注意，使用种子 1234 训练的模型在其 URL 中未指定种子。

使用多种种子复现较小 Pythia 模型的运行位于：https://wandb.ai/eleutherai/pythia-extra-seeds

勘误

如此 issue 所述，6.9B 和 12B 模型由于未在配置文件中指定初始化值，意外地使用了不同的初始化方式。

使用 Pythia

快速开始

所有 Pythia 模型都托管在 Huggingface hub 上。它们可以通过以下代码加载和使用（以 3000 步的 pythia-70M-deduped 模型检查点为例）：

from transformers import GPTNeoXForCausalLM, AutoTokenizer

model = GPTNeoXForCausalLM.from_pretrained(
  "EleutherAI/pythia-70m-deduped",
  revision="step3000",
  cache_dir="./pythia-70m-deduped/step3000",
)

tokenizer = AutoTokenizer.from_pretrained(
  "EleutherAI/pythia-70m-deduped",
  revision="step3000",
  cache_dir="./pythia-70m-deduped/step3000",
)

inputs = tokenizer("Hello, I am", return_tensors="pt")
tokens = model.generate(**inputs)
print(tokenizer.decode(tokens[0]))

所有模型均在批次大小为 2,097,152 个 token 的情况下训练了相当于 143000 步。版本/分支 step143000 与每个模型 main 分支上的模型检查点完全对应。

我们还有所有以 GPT-NeoX 库 接受的格式存在的模型检查点，最终步检查点+优化器状态可从 Hugging Face Hub 的 EleutherAI/neox-ckpt-pythia-xxx-deduped-v1 下载，但由于优化器状态的大小以及预期需求较低，我们不会大规模提供所有步的检查点。如果您想使用 GPT-NeoX 代码库中的中间模型进行分析，或需要其他步的优化器状态，请发送电子邮件至 hailey@eleuther.ai 和 stella@eleuther.ai。

❗ Huggingface 上大小为 160m, 410m, 1.4b 的 pythia-{size}-v0 模型是在批次大小为 4M token、共计 71500 步的条件下训练的，并且每 500 步保存一次检查点。为了与所有批次大小为 2M 的模型保持一致，这些 v0 模型在 Huggingface 上的步骤名称已被重命名。因此，标记为 step1000 的 pythia-1.4b-v0 模型检查点实际上是第 500 步，但它已经观察到了与其他 step1000 检查点相同数量的 token。

复现训练

(由 @BaruchG 提供的扩展复现说明)。

我们提供训练数据用于复现我们的训练运行。GPT-NeoX 库 需要将预分词后的训练数据以两个内存映射的 numpy 数组形式提供：一个 .bin 文件和一个 .idx 文件。我们通过 Hugging Face hub 提供这些文件。要下载并使用去重后的 Pile 训练数据：

git lfs clone https://huggingface.co/datasets/EleutherAI/pythia_deduped_pile_idxmaps

# 可选：检查文件完整性
python utils/checksum_shards.py

python utils/unshard_memmap.py --input_file ./pythia_deduped_pile_idxmaps/pile_0.87_deduped_text_document-00000-of-00082.bin --num_shards 83 --output_dir ./pythia_pile_idxmaps/

# 完整文件的正确 sha256 是 0cd548efd15974d5cca78f9baddbd59220ca675535dcfc0c350087c79f504693
# 可以使用 sha256sum ./pythia_pile_idxmaps/* 进行检查

这将需要超过一天的时间来运行，但应该不需要超过 5 GB 的 RAM。我们建议下载此文件，而不是从头开始重新标记化 Pile，以保证保留 Pythia 模型所见的数据顺序。除了训练数据，您还需要创建一个我们用于训练模型的 tokenizer 的本地副本。您可以在此处找到它。

接下来，您需要设置训练环境：

git clone https://github.com/EleutherAI/gpt-neox.git
cd gpt-neox
git checkout v1.0
pip install -r requirements/requirements-flashattention.txt
wget https://github.com/EleutherAI/pythia/blob/main/models/160M/pythia-160m-deduped.yml
docker build -t pythia:latest .

在容器构建完成后，使用以下命令运行容器（从 GPT-NeoX 仓库的根目录，并确保您的 pythia yaml 文件在该文件夹内可访问）：

docker run --runtime=nvidia --rm -it -e NVIDIA_VISIBLE_DEVICES=0,1,2,3 --shm-size=1g --ulimit memlock=-1 --mount type=bind,src=$PWD,dst=/gpt-neox -v $(pwd):/workspace/ pythia:latest bash

您可以使用 -v 参数添加更多挂载卷，以便数据集和 Yaml 文件在 docker 容器内可访问。

如下所示更改数据路径和 tokenizer 路径的行：

  "train-data-paths": ["/fsx/pile/pile_20B_tokenizer_text_document"], # 指向步骤 1 中生成的包含 .bin 和 .idx 文件的文件夹
  "valid-data-paths": ["/fsx/pile/pile_20B_tokenizer_text_document"], # 指向步骤 1 中生成的包含 .bin 和 .idx 文件的文件夹
  "test-data-paths": ["/fsx/pile/pile_20B_tokenizer_text_document"], # 指向步骤 1 中生成的包含 .bin 和 .idx 文件的文件夹

  "tokenizer-type": "HFTokenizer",
  "vocab-file": "/fsx/pile/20B_tokenizer.json", # 指向步骤 2 中获取的 tokenizer

根据您可用的 VRAM 大小，您可能需要调整批次大小。总批次大小的计算公式为 Total GPUs * train_micro_batch_size_per_gpu * gradient_accumulation_steps / (pipe-parallel-size * model-parallel-size)，并且需要保持为 1024 以匹配 Pythia 训练批次大小。

   "train_micro_batch_size_per_gpu": XXX, # 设置为适合您 GPU 内存的值
   "gradient_accumulation_steps": 1, # 调整此值以使总批次大小为 1024

如果您希望保存权重，请也将该信息添加到 yaml 文件中。例如，要保存在检查点文件夹中，可以在底部添加：

  "launcher": "slurm",
  "deepspeed_slurm": false,

  "save": "checkpoints",
  "load": "checkpoints",
  "checkpoint_validation_with_forward_pass": False,
}

确保路径是 docker 容器内部的路径。如果您希望权重持久化，请确保它们可以从容器外部访问，例如在 /workspace/ 中。

现在，您应该能够通过运行以下命令开始训练您的模型：

python deepy.py train.py pythia-160m-deduped.yml  2>&1 | tee output.txt

输出将保存到 output.txt，如果不需要，可以删除结尾部分。

要将您的模型转换为 Hugging Face transformers 格式，您可以使用 GPT-NeoX 库中的脚本 tools/convert_to_hf.py。您可能需要在文件类型中添加 from typing import List 并将此处的行从 list[torch.Tensor] 更改为 List[torch.Tensor]。然后，您可以像这样运行脚本以转换步骤 143000 的权重：

python tools/convert_to_hf.py --input_dir checkpoints/global_step143000/ --config_file checkpoints2/global_step 143000/configs/pythia-70m.yml --output_dir ./output/

这应该会输出一个与 https://huggingface.co/EleutherAI/pythia-70m-deduped/tree/main 类似的文件结构。

❗ 有时，人们会发现由于我们无法调试的原因，他们最终没有获得正确的 tokenizer。如果您的 tokenizer_config.json 看起来与此处的不同，并且 special_tokens_map.json 看起来与此处的不同，您可能需要用 Huggingface 上的文件替换它们。

要使用我们的评估库运行评估，请安装此处的容器（已在 4.28 和 4.29 版本上测试）。设置好该 docker 容器后，运行：

git clone https://github.com/EleutherAI/lm-evaluation-harness
cd lm-evaluation-harness
pip install -e .

如 Harness 仓库中所述。然后，您应该能够通过指向您的权重（应在容器中）来运行基准测试，运行类似以下的命令：

python3 main.py --model hf-causal-experimental  --model_args pretrained=../gpt-neox/output/ --tasks lambada_openai,piqa,winogrande,arc_easy,sciq,wikitext --device cuda:0

探索数据集

我们提供一个工具来查看训练期间所有模型使用的特定部分训练数据加载器，位于 utils/batch_viewer.py。

首先，我们需要克隆 Pythia 仓库：

git clone https://github.com/EleutherAI/pythia

接下来，我们必须安装依赖项：

pip install torch==1.13.0+cu117 -f https://download.pytorch.org/whl/torch/
pip install numpy tqdm huggingface_hub

接下来，我们必须下载相应的数据集。我们提供重复和去重 Pile 的预打乱版本。使用 Huggingface 的实用程序下载相应版本，如下所示：

提示：确保将 path/to/* 替换为您打算保存从 Huggingface 下载的数据集的适当路径。
- 要下载标准版本，请使用
py from huggingface_hub import hf_hub_download hf_hub_download(repo_id="EleutherAI/pile-standard-pythia-preshuffled", repo_type="dataset", cache_dir="path/to/local/folder")
- 要下载去重版本，请使用
py from huggingface_hub import hf_hub_download hf_hub_download(repo_id="EleutherAI/pile-deduped-pythia-preshuffled", repo_type="dataset", cache_dir="path/to/local/folder")

现在，您可以使用脚本 utils/unshard_memmap.py 合并文件：

python3 utils/unshard_memmap.py --input_file "path/to/local/folder/document-00000-of-00020.bin" --num_shards 21 --output_dir "path/to/merged/folder/"

确保也将索引文件复制到合并后的文件夹中，使用命令：

cp path/to/local/folder/document.idx path/to/merged/folder/document.idx

现在，我们已经准备好运行 utils/batch_viewer.py！

python3 utils/batch_viewer.py \
  --start_iteration 0 \
  --end_iteration 1000 \
  --load_path path/to/merged/folder/document \
  --save_path path/to/save/folder/ \
  --conf_dir utils/dummy_config.yml

这会将包含所有索引的单独文件保存为 numpy 数组。

现在，您可以使用 numpy 加载它，如下所示：

import numpy as np

indicies = np.load("path/to/save/folder/indicies.npy")

这些索引包含大小为 (None, 2049) 的标记化整数序列，其中每个整数对应一个唯一的 token 索引。请注意，文档被连接起来，并由一个 EOD token 分隔。因此，每个样本或批次可能不以 EOD token 开头。在训练期间，目标 token 会左移 1 位。因此，序列长度为 2048 的模型需要长度为 2049 的序列进行训练（有关更多信息，请参阅此评论）。

Pythia 论文复现

我们为本仓库的 case-studies/ 文件夹中对复现 Pythia 套件论文中进行的案例研究感兴趣的人提供了更多信息。

基准测试分数

我们还提供了在各种 NLP 数据集上的 0-shot 和 5-shot 基准测试结果：

• ARC-challenge (arc_challenge)
• ARC-easy (arc_easy)
• BLiMP (blimp_*)
• Lambada (lambada_openai)
• LogiQA (logiqa)
• MMLU (hendrycksTest*)
• PiQA (piqa)
• SciQ (sciq)
• Wikitext (wikitext)
• Winogrande (winogrande)
• WSC (wsc)

评估是使用 LM Evaluation Harness 在 GPT-NeoX 中执行的，可以通过模型和步骤在本仓库的 evals/pythia-v1/*/* 中查看。警告： 所有评估都是在大约一年前使用语言模型评估 harness 的 to-do 提交运行的，可能无法由当前版本复现。

基于 Pythia 的研究

我们 Pythia 项目的主要目标是支持 EleutherAI 内以及整个社区在可解释性和学习动态等主题上的研究。我们在此记录了一些使用我们模型的精选论文，重点关注那些因 Pythia 套件而具有独特优势，并且使用其他组织发布的模型难以实现或无法实现的工作。有关引用 Pythia 的更大列表，请参阅此处。

语言模型内部机制

• Belrose, et al. "Eliciting latent predictions from transformers with the tuned lens." arXiv preprint arXiv:2303.08112 (2023). EleutherAI 论文
• Brown, et al. "Understanding the Inner Workings of Language Models Through Representation Dissimilarity." Conference on Empirical Methods in Natural Language Processing (2023).
• Feng and Steinhardt. "How do Language Models Bind Entities in Context?." International Conference on Learning Representations (2023).
• Garde, Kran, and Barez. "DeepDecipher: Accessing and Investigating Neuron Activation in Large Language Models." arXiv preprint arXiv:2310.01870 (2023).
• Gurnee, et al. "Finding Neurons in a Haystack: Case Studies with Sparse Probing." Transactions of Machine Learning Research (2023).
• Stolfo, Belinkov, and Sachan. "Understanding Arithmetic Reasoning in Language Models using Causal Mediation Analysis." Conference on Empirical Methods in Natural Language Processing (2023).

学习动态

• Gupta, et al. "Continual Pre-Training of Large Language Models: How to re-warm your model?." Workshop on Efficient Systems for Foundation Models @ ICML (2023).
• Michaelov and Bergen. "Emergent inabilities? Inverse scaling over the course of pretraining." Findings of the Association for Computational Linguistics: EMNLP (2023).
• Sanyal, et al. "Understanding the Effectiveness of Early Weight Averaging for Training Large Language Models." arXiv preprint arXiv:2306.03241 (2023).
• Tian, et al. "JoMA: Demystifying Multilayer Transformers via JOint Dynamics of MLP and Attention." arXiv preprint arXiv:2310.0053 (2023).
• Ye, et al. "Language Versatilists vs. Specialists: An Empirical Revisiting on Multilingual Transfer Ability." arXiv preprint arXiv:2306.06688 (2023).
• Belrose, et al. "Neural Networks Learn Statistics of Increasing Complexity." International Conference on Learning Representations (2024). EleutherAI 论文
• Godey et al. "Why do small language models underperform? Studying Language Model Saturation via the Softmax Bottleneck." arXiv preprint arXiv:2404.07647 (2024).
• Singh, et al. "Hallmarks of Optimization Trajectories in Neural Networks: Directional Exploration and Redundancy." arXiv preprint arXiv:2403.07379 (2024).
• Tigges, et al. "Stability and Generalizability of Language Model Mechanisms Across Training and Scale." Mechanistic Interpretability Workshop @ ICML (2024). EleutherAI 论文
• Diehl Martinez, et al. "Tending Towards Stability: Convergence Challenges in Small Language Models " Findings of the Association for Computational Linguistics: EMNLP (2024).

训练数据如何决定模型行为

• Roger. "Large Language Models Sometimes Generate Purely Negatively-Reinforced Text." arXiv preprint arXiv:2306.07567 (2023).
• Oh, et al. "Frequency Explains the Inverse Correlation of Large Language Models’ Size, Training Data Amount, and Surprisal’s Fit to Reading Times." arXiv preprint arXiv:2402.02255 (2024).
• Liu, et al. "On Training Data Influence of GPT Models." arXiv preprint arXiv:2404.07840 (2024).
• Lesci, et al. "Causal Estimation of Memorisation Profiles." Association for Computational Linguistics (2024).

安全、审计与合规性研究

• Ippolito, et al. "Reverse-Engineering Decoding Strategies Given Blackbox Access to a Language Generation System." International Natural Language Generation Conference. 2023.
• Biderman, et al. "Emergent and predictable memorization in large language models." Neural Information Processing Systems (2023). EleutherAI 论文
• Choi, Shavit, and Duvenaud. "Tools for Verifying Neural Models' Training Data." Neural Information Processing Systems (2023).
• Li, et al. "MoPe: Model Perturbation-based Privacy Attacks on Language Models." Conference on Empirical Methods in Natural Language Processing (2023).
• Min, et al. "SILO Language Models: Isolating Legal Risk In a Nonparametric Datastore." International Conference on Learning Representations (2024).
• Pawelczyk, et al. "Machine Unlearning Fails to Remove Data Poisoning Attacks." arXiv preprint arXiv:2406.17216 (2024).
• Prashanth, et al. "Recite, Reconstruct, Recollect: Memorization in LMs as a Multifaceted Phenomenon." arXiv preprint arXiv:2406.17746 (2024). EleutherAI 论文
• Duan, et al. "Do Membership Inference Attacks Work on Large Language Models?." Conference on Language Modeling (2024).

引用详情

如果您在研究中使用 Pythia 模型，请引用我们的论文：

@inproceedings{biderman2023pythia,
  title={Pythia: A suite for analyzing large language models across training and scaling},
  author={Biderman, Stella and Schoelkopf, Hailey and Anthony, Quentin Gregory and Bradley, Herbie and O’Brien, Kyle and Hallahan, Eric and Khan, Mohammad Aflah and Purohit, Shivanshu and Prashanth, USVSN Sai and Raff, Edward and others},
  booktitle={International Conference on Machine Learning},
  pages={2397--2430},
  year={2023},
  organization={PMLR}
}

如果您使用本仓库中其他论文的数据或结果，请引用相应的论文。引用信息可在各自的 README 中找到，为了方便起见，也转载如下：

```
@inproceedings{biderman2023emergent,
title={Emergent and Predictable Memorization in Large Language Models},
author={Biderman, Stella and Prashanth, USVSN Sai and Sutawika, Lintang and Schoelkopf, Hailey and Anthony, Quentin and Purohit, Shivanshu and Raff, Edward},
booktitle={Advances in Neural Information Processing Systems},
year={2023}
}

@inproceedings{van2025polypythias,
title={PolyPythias: Stability and Outliers across Fifty Language Model Pre-Training Runs},
author={van der