椰子

该代码库是 在连续潜在空间中训练大型语言模型进行推理 的官方实现。

快速入门

克隆仓库：

git clone git@github.com:facebookresearch/coconut.git
cd coconut

设置环境：

conda create --name coconut python=3.12
conda activate coconut
pip install -r requirements.txt

代码依赖 wandb 进行日志记录。请在运行任何实验之前，按照此 文档 登录您的 wandb 账户。

数据

用于训练和评估的数据应以如下所示的 JSON 文件形式呈现：

[
  {
    "question": "...",
    "answer": "...",
    "steps": ["...", "...", ...]
  },
  ...
]

该文件应包含一个数据点列表。每个数据点由一个问题（字符串）、一个答案（字符串）和步骤列表（字符串）组成，其中每一项都是字符串。

例如，您可以通过运行以下命令下载并处理 GSM8K 数据集（带有 增强的训练和验证集）：

bash preprocessing/gsm_icot.bash

参数

一次运行的配置应在 YAML 文件中指定（示例可参见 这里）。

• 通用设置

  • project: wandb 的项目名称
  • save_path: 您存储检查点的路径
  • only_eval: 如果为真，则仅加载模型并在 val_path 中的数据上进行测试（必须与 load_model_path 一起使用）。否则，在 train_path 上训练模型，并在每个 epoch 结束后在 val_path 上进行测试。

• 方法

  • coconut: 训练椰子模型
  • cot: 训练 CoT 模型
  • no_thoughts: 训练无思维的椰子模型
  • no_cot: 训练无 CoT 模型

• 训练设置

  • c_thought: 每个推理步骤中的连续思维数量
  • epochs_per_stage: 每个训练阶段的 epoch 数
  • max_latent_stage: 最大训练阶段数（除初始阶段外）
  • pad_latent_to_max: 如果推理步骤的数量少于当前训练阶段的索引，则填充连续思维的数量。
  • save_only_improve: 仅当最佳验证准确率更新时才保存模型。建议在椰子模型训练中将其设置为 False，否则最后阶段的检查点可能不会被保存。
  • uniform_prob: 混合其他阶段数据的概率。标准实验为 0，分析实验为 0.3。
  • model_id: 用于初始化的 Huggingface 模型 ID，例如 openai-community/gpt2
  • load_model_path: 要加载的检查点路径。用于两种情况：(1) 用于评估 (2) 从 CoT 微调模型初始化椰子模型。
  • seed: 随机种子。
  • resume: 继续训练的 epoch。可用于跳过初始训练阶段。
  • bf16: 是否使用 bf16 训练。
  • train_path: 训练集路径。
  • val_path: 验证或测试集路径（取决于 only_eval)
  • reset_optimizer: 切换训练阶段时是否重置优化器。
  • batch_size_training: 每块 GPU 上用于训练模型的批量大小。
  • debug: 如果为真，则不进行 wandb 日志记录和模型保存。将使用数据的一个子集。
  • gradient_accumulation_steps: 梯度累积步数
  • num_epochs: 最大训练 epoch 数。
  • lr: 学习率
  • weight_decay: 权重衰减

训练

运行以下命令（替换 N_GPUS 和 PATH_TO_ARGS）：

torchrun --nnodes 1 --nproc_per_node N_GPUS run.py PATH_TO_ARGS

复现实验

在此我们提供复现论文中实验的说明。

以下所有命令均假设使用 4 块 A100（80GB）GPU。您可以根据自己的资源情况，在启动运行时更改配置文件中的相应参数（batch_size_training、gradient_accumulation_steps）以及 nproc_per_node。

GSM8K

预处理数据：

bash preprocessing/gsm_icot.bash

首先以 CoT 方式训练模型（作为第 0 阶段训练）

torchrun --nnodes 1 --nproc_per_node 4 run.py args/gsm_cot.yaml

选择一个检查点作为椰子模型的初始化（预计验证准确率约为 40%）。将 args/gsm_coconut.yaml 中的 load_model_path 替换为您选择的检查点，然后运行：

torchrun --nnodes 1 --nproc_per_node 4 run.py args/gsm_coconut.yaml

找到验证准确率最高的检查点，并将其路径放入 args/gsm_coconut_eval.yaml 中的 load_model_path。进行评估：

torchrun --nnodes 1 --nproc_per_node 4 run.py args/gsm_coconut_eval.yaml

ProntoQA

请克隆 ProntoQA 的官方 GitHub 仓库，并使用以下命令生成原始数据集：

cd prontoqa
python run_experiment.py --model-name json --model-size dummy --ordering random --num-trials 10000 --few-shot-examples 0 --ontology fictional --min-hops 5 --max-hops 5 --hops-skip 1

然后将生成的 5hop_0shot_random.json 文件复制到 data 目录，并使用以下命令预处理数据集：

python preprocessing/prontoqa.py

随后运行以下命令训练模型：

torchrun --nnodes 1 --nproc_per_node 4 run.py args/prontoqa_coconut.yaml

找到验证准确率最高的检查点，并将其路径放入 args/prosqa_coconut_eval.yaml 中的 load_model_path。进行评估：

torchrun --nnodes 1 --nproc_per_node 4 run.py args/prosqa_coconut_eval.yaml

ProsQA

ProsQA 数据集位于 data/prosqa_*.json。

随后运行以下命令训练模型：

torchrun --nnodes 1 --nproc_per_node 4 run.py args/prosqa_coconut.yaml

找到验证准确率最高的检查点，并将其路径放入 args/prosqa_coconut_eval.yaml 中的 load_model_path。进行评估：

torchrun --nnodes 1 --nproc_per_node 4 run.py args/prosqa_coconut_eval.yaml

引用

如果您在研究中使用此代码库，请使用以下 BibTex 条目引用我们的论文：

@article{hao2024training,
  title={Training Large Language Models to Reason in a Continuous Latent Space},
  author={Hao, Shibo and Sukhbaatar, Sainbayar and Su, DiJia and Li, Xian and Hu, Zhiting and Weston, Jason and Tian, Yuandong},
  journal={arXiv preprint arXiv:2412.06769},
  year={2024}
}

许可证

此代码以 MIT 许可证发布（参见 LICENSE）。

Coconut 快速上手指南

Coconut 是论文《Training Large Language Models to Reason in a Continuous Latent Space》的官方实现，旨在训练大语言模型在连续潜在空间中进行推理。

环境准备

• 操作系统: Linux (推荐)
• Python: 3.12
• GPU: 支持 CUDA 的 NVIDIA GPU (复现实验建议配备多卡，如 4x A100)
• 依赖管理: Conda
• 外部服务: 需注册 Weights & Biases (wandb) 账号用于实验日志记录

安装步骤

1. 克隆代码库

git clone git@github.com:facebookresearch/coconut.git
cd coconut

2. 创建并激活虚拟环境

conda create --name coconut python=3.12
conda activate coconut

3. 安装依赖

建议使用国内镜像源加速安装（如清华源）：

pip install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple

4. 配置 Wandb

在运行任何实验前，请登录您的 wandb 账号：

wandb login

(按提示输入 API Key，若无法访问官网，可在国内网络环境下操作或使用代理)

基本使用

1. 数据准备

训练数据需为 JSON 格式，包含 question、answer 和 steps 字段。 以 GSM8K 数据集为例，运行以下脚本进行预处理：

bash preprocessing/gsm_icot.bash

2. 配置文件

修改 args/ 目录下的 YAML 配置文件（如 args/gsm_coconut.yaml）。 关键参数说明：

• train_path: 训练集路径
• val_path: 验证集路径
• save_path: 模型检查点保存路径
• model_id: 初始化的 HuggingFace 模型 ID (如 openai-community/gpt2)
• load_model_path: 若从 CoT 微调后的模型继续训练，在此填入检查点路径

3. 启动训练

使用 torchrun 启动训练。以下命令示例使用 4 张 GPU：

torchrun --nnodes 1 --nproc_per_node 4 run.py args/gsm_coconut.yaml

典型两阶段训练流程（参考 GSM8K 复现）：

1. 阶段一：训练 CoT 模型

   torchrun --nnodes 1 --nproc_per_node 4 run.py args/gsm_cot.yaml
   

2. 阶段二：训练 Coconut 模型
   • 选择阶段一中验证准确率较高（约 40%）的检查点。
   • 在 args/gsm_coconut.yaml 中设置 load_model_path 为该检查点路径。
   • 运行：

   torchrun --nnodes 1 --nproc_per_node 4 run.py args/gsm_coconut.yaml
   

4. 模型评估

找到验证集表现最佳的检查点，在评估配置文件（如 args/gsm_coconut_eval.yaml）中设置 load_model_path，并将 only_eval 设为 true，然后运行：

torchrun --nnodes 1 --nproc_per_node 4 run.py args/gsm_coconut_eval.yaml