Huginn-0125 的预训练与推理代码——一种深度递归模型

本仓库包含我们在 Frontier 超级计算机上，使用 4096 张 AMD GPU 大规模训练递归深度模型所用的代码。关于该模型的所有细节均可在技术报告《通过潜在推理扩展测试时计算：一种递归深度方法》中找到（https://www.arxiv.org/abs/2502.05171）。最终模型名为 huginn-0125，可在以下链接获取：https://huggingface.co/tomg-group-umd/huginn-0125。

随着时间推移，我们还积累了不少用于测试和推理的代码，尤其是与 Hugging Face 兼容的建模文件 recpre/raven_modeling_minimal.py 以及位于 evaluate_raven 中的推理脚本。支持通过 lm-eval 进行基准测试，更多信息见下文。

本仓库最初基于 https://github.com/Lightning-AI/litgpt 的一个分支，该分支对我们起步阶段的工作帮助很大，但目前几乎已不保留 litgpt 的代码。本仓库中的代码由 Jonas Geiping、John Kirchenbauer、Sean McLeish、Khalid Saifullah、Manli Shu、Neel Jain、Siddarth Singh、Abhimanyu Hans、Monte Hoover 和 Prajwal Singhanaia 共同编写。

此外，本仓库还包含了我们用于准备模型分词器及数据的所有代码，主要位于 scripts/ 目录下。

我个人并不一定建议您使用此实现来预训练自己的模型，但我希望它能作为一份有用的参考，展示我们为运行该模型所做出的具体选择，以及在 AMD 系统的限制下如何实际运行该模型。如果您正在使用其中任何一项，请随时提出问题以获取更多细节。

代码设置：

• 实际的模型定义位于 recpre/model_dynamic.py。
• 训练流程由 train.py 协调。
• 模型结构信息可在 recpre/model_registry.py 中找到。最终模型的结构为 nebel-raven-3.5b。
• 我们的两次大规模运行配置位于 launch_configs/ 目录中。
• 环境标志可从 launch_frontier.py 中读取。
• 并行化实现深藏于 recpre/utils.py 中的一个名为 SimpleFabric 的类中。在撰写本文时，为了在使用 OFI 插件的大规模场景下解决 RCCL 死锁问题，我们采用了 _allreduce_chunk_stream 方法进行节点间通信，这是唯一的解决方案。

如果您只想了解模型架构，那么用于推理的代码可能更容易阅读。这些代码可以在该模型的所有 Hugging Face 仓库中找到，也可以在 recpre/raven_modeling_minimal.py 中查看。

重现基准测试分数

论文中报告的所有基准测试分数均使用 lm-eval 测试框架计算得出，除了代码任务，后者是使用 bigcode 执行的。对于默认基准测试，您可以按如下方式运行 lm-eval（无需安装）：

lm_eval --model hf --model_args pretrained=tomg-group-umd/huginn-0125,trust_remote_code=True,dtype=bfloat16,mean_recurrence=32 --tasks hellaswag --batch_size=auto --num_fewshot=0

对于 GSM8k 任务，“带系统提示”的版本是指使用以下系统提示和聊天格式的调用方式：

lm_eval --model hf  \
--model_args pretrained=tomg-group-umd/huginn-0125,trust_remote_code=True,dtype=bfloat16,mean_recurrence=32  \ 
--tasks gsm8k_cot  --batch_size=auto  --apply_chat_template=True --fewshot_as_multiturn \
--system_instruction="You are a helpful assistant that can assist users with mathematical reasoning."  \ 

要重现 humaneval 分数，如今无需直接安装 bigcode-eval，也可以使用 lm-eval 测试框架，具体命令如下：

HF_ALLOW_CODE_EVAL=1 accelerate launch -m lm_eval \
--model hf --model_args pretrained=tomg-group-umd/huginn-0125,mean_recurrence=32,trust_remote_code=True,dtype=bfloat16  \
--tasks humaneval_instruct --batch_size=1 --num_fewshot=0  \
--output_path=outputs/heval --confirm_run_unsafe_code \
--apply_chat_template=True \
 --gen_kwargs=do_sample=True,temperature=0.2,top_p=0.95

快速推理

现已支持通过 vllm 进行快速推理。更多详情请参阅 vllm 文件夹。适用于该模型的插件可安装到任何较新的 vllm v1 版本中。

数据

我们已将整个训练数据集上传至 Hugging Face，您可在此处找到：https://huggingface.co/datasets/tomg-group-umd/huginn-dataset。该上传包含 4096 个 Parquet 文件（分别用于训练和验证），正好对应用于数据并行训练的 4096 张 AMD GPU。每行包含 4096+1 个 token，因此一次训练步骤（我们采用局部微批次大小为 1）恰好对应每个文件的一行。

细节说明

如果您想在 AMD 集群上复现本次模型训练和数据收集流程，需要采取哪些步骤呢？请按照以下步骤操作（遇到困难时请随时咨询更多细节）：

1. 使用 scripts/tokenizer_generation.py 生成分词器。运行脚本前，请根据您的系统调整所有路径。数据下载是自动进行的。您还需要来自 https://github.com/gautierdag/bpeasy 的 BPE 训练工具。
2. 运行 scripts/scalable_data_download.py 下载所有原始数据集。该脚本名称名不副实，实际上并不具备真正的可扩展性，耗时较长且会占用大量存储空间，并可能因随机错误而失败。您还会发现脚本中针对一些格式不良的数据集硬编码了若干规则。在您运行此脚本时，某些数据集作者可能已经更新了他们的数据集，从而打破了此处设定的假设。在这种情况下，您可能会遇到错误，需要进一步调查相关数据集。完成此步骤后，所有原始数据集将被放置在 staging 文件夹中。
3. 运行 scripts/parquet_to_parquet_tokenizer.py 生成分词后的数据集。同样，请务必正确设置路径。
4. 分词完成后，运行 scripts/parquet_to_parquet_shuffler.py 对数据进行打乱。
5. 在 launch_configs/ 目录中定义您自己的启动配置，或直接使用我们的配置，然后在您的集群上启动 train.py。我们是在 Frontier 超级计算机上使用名为 launch_frontier.py 的启动脚本进行训练的，但这对其他集群并不适用。设置大规模运行时，请遵循您所在集群的最佳实践和环境标志指南。核心命令仅为 python train.py --config=launch_configs/your_config.yaml。
6. 观察训练过程（希望一切顺利）。您可根据需要向训练过程中添加额外的环节。

许可证

本代码采用 apache-2.0 许可证发布。部分代码则采用 Lightning AI 的 Apache-2.0 许可证。

引用

@article{geiping_scaling_2025,
  title = {利用{{潜在推理}}扩展{{测试时计算}}：{{一种递归深度方法}}},
  shorttitle = {利用{{潜在推理}}扩展{{测试时计算}}},
  author = {Geiping, Jonas 和 McLeish, Sean 和 Jain, Neel 和 Kirchenbauer, John 和 Singh, Siddharth 和 Bartoldson, Brian R. 和 Kailkhura, Bhavya 和 Bhatele, Abhinav 和 Goldstein, Tom},
  year = {2025},
  month = feb,
  eprint = {2502.05171},
  primaryclass = {cs},
  publisher = {arXiv},
  doi = {10.48550/arXiv.2502.05171},
  url = {http://arxiv.org/abs/2502.05171},
  urldate = {2025-02-10},
  archiveprefix = {arXiv},
  keywords = {计算机科学 - 计算与语言,计算机科学 - 机器学习},
  journal = {arxiv:2502.05171[cs]}
}

联系方式

如有任何问题，请随时联系我们，或提交一个问题！

recurrent-pretraining 快速上手指南

本指南旨在帮助开发者快速了解并使用 recurrent-pretraining 项目，该项目包含了深度循环模型 Huginn-0125 的预训练、推理及评估代码。该模型在 4096 块 AMD GPU (Frontier) 上训练而成，专注于通过潜在推理扩展测试时计算能力。

环境准备

系统要求

• 硬件: 推荐拥有多卡环境的服务器（原项目在 AMD MI250X/GPU Frontier 集群运行）。支持 NVIDIA GPU，但针对 AMD 系统有特定的通信优化。
• 操作系统: Linux (Ubuntu/CentOS 等)。
• Python: 建议 Python 3.10+。
• 深度学习框架: PyTorch (需适配对应 GPU 架构)，Lightning AI (部分基础代码源自 litgpt)。

前置依赖

• 推理/评估: 无需完整安装训练环境，仅需 transformers, accelerate, lm-eval-harness。
• 训练复现: 需要完整的源码环境，以及用于分词训练的 bpeasy 工具。
• 网络: 访问 Hugging Face 下载模型和数据（国内用户建议配置镜像源或使用代理）。

安装步骤

1. 克隆仓库

git clone https://github.com/tomg-group-umd/recurrent-pretraining.git
cd recurrent-pretraining

2. 安装依赖

由于项目基于 Lightning AI fork 修改，建议先安装基础 PyTorch 环境（根据你的 GPU 类型选择 CUDA 或 ROCm 版本），然后安装项目依赖。

# 示例：安装基础依赖 (具体 requirements.txt 若存在请优先使用)
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 或对应的 rocm 版本
pip install accelerate lm-eval transformers sentencepiece

注意: 如果要复现完整的数据预处理和训练流程，还需手动安装 BPE 分词器：

git clone https://github.com/gautierdag/bpeasy.git
cd bpeasy && pip install .

3. 验证模型可用性 (可选)

确保可以从 Hugging Face 拉取模型。国内用户如遇到连接超时，可设置环境变量使用镜像：

export HF_ENDPOINT=https://hf-mirror.com

基本使用

本项目最核心的用途是推理和基准测试。以下是使用 lm-eval harness 对已发布的 huginn-0125 模型进行评测的最简示例。

场景一：运行标准基准测试 (Hellaswag)

无需额外安装，直接通过命令行调用 lm-eval 加载模型并进行测试。

lm_eval --model hf \
--model_args pretrained=tomg-group-umd/huginn-0125,trust_remote_code=True,dtype=bfloat16,mean_recurrence=32 \
--tasks hellaswag \
--batch_size=auto \
--num_fewshot=0

场景二：运行数学推理任务 (GSM8k)

包含系统提示词（System Prompt）和聊天模板格式的调用示例：

lm_eval --model hf \
--model_args pretrained=tomg-group-umd/huginn-0125,trust_remote_code=True,dtype=bfloat16,mean_recurrence=32 \
--tasks gsm8k_cot \
--batch_size=auto \
--apply_chat_template=True \
--fewshot_as_multiturn \
--system_instruction="You are a helpful assistant that can assist users with mathematical reasoning."

场景三：运行代码生成任务 (HumanEval)

需要允许执行代码并设置特定的生成参数：

HF_ALLOW_CODE_EVAL=1 accelerate launch -m lm_eval \
--model hf \
--model_args pretrained=tomg-group-umd/huginn-0125,mean_recurrence=32,trust_remote_code=True,dtype=bfloat16 \
--tasks humaneval_instruct \
--batch_size=1 \
--num_fewshot=0 \
--output_path=outputs/heval \
--confirm_run_unsafe_code \
--apply_chat_template=True \
--gen_kwargs=do_sample=True,temperature=0.2,top_p=0.95

进阶：查看模型架构

如果仅想研究模型结构而不进行大规模评测，可直接查看精简版的建模文件：

• 路径：recpre/raven_modeling_minimal.py
• 说明：此文件兼容 Hugging Face 格式，适合快速阅读核心递归深度（Recurrent Depth）实现逻辑。

进阶：本地训练复现 (简要)

若需在集群上复现训练过程，请按以下顺序执行脚本（需自行配置 launch_configs 中的路径和参数）：

1. 生成分词器：python scripts/tokenizer_generation.py
2. 下载原始数据：python scripts/scalable_data_download.py
3. 数据分词：python scripts/parquet_to_parquet_tokenizer.py
4. 数据打乱：python scripts/parquet_to_parquet_shuffler.py
5. 启动训练：python train.py --config=launch_configs/your_config.yaml