跳到主要内容

数据截至 (上游 commit 1416fa0cf215)

01 · 复现路线全景：从 R1 数据到蒸馏模型

这一章讲什么： DeepSeek-R1 论文的训练管线被拆成了哪三步、各自验收标准是什么；然后沿着唯一已完成的「蒸馏路线」走一遍——R1 推理轨迹怎么批量生成、数据怎么混配、OpenR1-Distill-7B 那份 SFT 配方里每个关键参数在解决什么问题。

1. 它要解决的小问题​

DeepSeek-R1 论文给了结果和路线，没给代码、没给数据、没给超参数。「复现 R1」因此不是一个任务，而是一串任务：先得造出训练数据，再谈训练。

Open R1 的应对是把大目标切成三步可独立验收的子目标（README.md:33-36）：

步骤目标验收物状态（as-of 本 commit）
Step 1复现 R1-Distill：蒸馏 R1 的高质量语料训小模型Mixture-of-Thoughts 数据集 + OpenR1-Distill-7B已完成（README.md:44）
Step 2复现 R1-Zero：纯 RL 从基座练起数学/推理/代码的大规模可验证数据集 + GRPO 管线进行中（管线已就绪）
Step 3基座模型 → 多阶段训练 → RL 调优的全链路端到端配方进行中

这个拆法本身就值得学：每一步都产出一个独立可用的东西（数据集、管线、配方），而不是憋一个大 release。

2. 直觉：为什么「蒸馏」先行​

三步里先做的是蒸馏，原因很务实——它是唯一不需要解决开放问题的一步：

  • 教师模型现成的（DeepSeek-R1 开放权重）；
  • SFT 训练是成熟技术，trl 的 SFTTrainer 直接用；
  • 验收标准明确：分数对齐 DeepSeek 官方的 R1-Distill-Qwen-7B。

所以 Step 1 的全部难点压缩成一个问题：怎么造出几十万条「对的」长推理轨迹。这也正是本仓库数据侧代码（生成、去污染、过滤）比训练侧厚得多的原因。

3. 数据生成：671B 教师怎么批量「写解题过程」​

3.1 图示：生成作业的集群拓扑​

slurm/generate.slurm 把一次数据生成编成两个 Ray job：

Slurm 作业（2 节点 × 8 H100）
  │
  ├─► Ray job "vllm-server"（generate.slurm:176-187）
  │     vllm serve DeepSeek-R1
  │     TP=8(每节点) × PP=2(节点数)，OpenAI 兼容 API 暴露在头节点 :8000
  │
  └─► Ray job "generate"（generate.slurm:218-238）
        python generate.py --model deepseek-ai/DeepSeek-R1 ...
        distilabel 管线逐题打 HTTP 请求 → 每题 N 条轨迹 → push_to_hub

3.2 原理演示​

生成的核心就一件事：题目进来，包上固定 prompt 模板，请教师模型逐步推理并给出 \boxed{} 答案，一题采多条。

# 示意，非源码
prompt = "You will be given a problem. Please reason step by step, \
and put your final answer within \\boxed{}:\n" + problem

for problem in dataset:
    # 一题采 4 条，温度 0.6——采样多样性换来后续可筛选的空间
    traces = vllm_server.generate(prompt, n=4, temperature=0.6, max_tokens=8192)
    save(problem, traces)

3.3 真实实现​

管线本体在 build_distilabel_pipeline（src/open_r1/generate.py:23-63），结构简单得意外：

  • Pipeline().ray()（generate.py:45）——distilabel 管线跑在 Ray 上，客户端可以多副本并行打请求；
  • LLM 用的是 OpenAILLM（generate.py:47-54）——不直接嵌 vLLM，而是打一个已经起好的 OpenAI 兼容 server，客户端和 671B 推理服务彻底解耦；
  • num_generations + group_generations=True（generate.py:58-59）——一题多采样且按题分组落盘，这正是后面 pass rate 过滤（第 3 章）需要的数据形状。

一个容易错过的工程细节：api_key="something"（generate.py:49）——本地 vLLM server 不校验 key，但不能留空，这是 OpenAI 客户端的硬性要求。

3.4 坑：MoE 的 CUDA graph bug​

README 记录（README.md:733-738）：给 R1 做生成时需要一个 vLLM dev wheel，专门修 R1 的 CUDA graph capture 问题（MoE 模型的 moe_align_sum_kernels）。大模型蒸馏管线的真实瓶颈常常不是算法而是这种引擎级 bug——这也是为什么 generate.slurm 里 vLLM 要开 --enable-chunked-prefill、显存利用率压到 0.85（generate.slurm:180-187）。

4. 数据混配：多个数据集怎么合成一个训练集​

4.1 思路​

Step 1 的最终产物 Mixture-of-Thoughts 覆盖数学、代码、科学三类任务。训练侧需要一个机制：把几个 Hub 数据集按不同比例混成一个。Open R1 的做法是把「混配方案」也写进 YAML，让配置自描述。

4.2 真实实现​

配置层：ScriptArguments 在 trl 基类上加 dataset_mixture 字段（src/open_r1/configs.py:73-76），并在 __post_init__ 里做两件校验（configs.py:78-120）：

  • dataset_name 和 dataset_mixture 至少给一个（configs.py:79-80）；
  • 混配里所有数据集的 columns 列表必须一致（configs.py:113-120）——提前炸掉「两个数据集列名对不上、concat 出 NaN」这种静默错误。

执行层：get_dataset（src/open_r1/utils/data.py:12-65）的逻辑是——逐数据集加载 → 按 columns 裁剪 → 按 weight 等比子采样（shuffle(seed).select(range(int(len * weight)))，data.py:39）→ concatenate_datasets → 整体再 shuffle 一次（data.py:48）。注意 weight 是相对自身长度的保留比例，不是归一化的混合权重。

5. SFT 蒸馏配方逐项拆解​

OpenR1-Distill-7B 的完整配方在 recipes/OpenR1-Distill-7B/sft/config_distill.yaml。训练脚本本身（src/open_r1/sft.py:55 的 main）是标准的 trl 流程——读配置、加载、SFTTrainer 一把梭（sft.py:101-109）。干货全在 YAML 里。

5.1 关键旋钮一览​

参数值在解决什么
model_name_or_pathopen-r1/Qwen2.5-Math-7B-RoPE-300k（:3）不用原版 Qwen2.5-Math-7B，而用 RoPE 外推到 300k 上下文的改造版——推理轨迹动辄几万 token，基座的 4k 上下文装不下
max_length32768（:34）单条样本允许 32k token，配合上面的长上下文基座
learning_rate4.0e-05（:25）比常规 SFT 高（常见 1~2e-5）；推理蒸馏要「覆盖」基座原有的回答风格
num_train_epochs5（:36）5 个 epoch，蒸馏语料要反复吃
per_device_train_batch_size2（:40）+ 梯度累积 8（:19）32k 长度下每张 H100 只放得下 2 条；全局 batch = 2×8×8卡=128
use_liger_kerneltrue（:47）Liger 融合 kernel，长序列 SFT 的显存救星
eos_token<|im_end|>（:13）见 §5.2 的坑

5.2 坑：EOS 与 chat template 必须对齐​

README 用 diff 格式专门警告（README.md:158-175）：Qwen 系基座模型自带 chat template，训练时 EOS 必须显式设成 <|im_end|>，否则模型学不会「停」。

代码侧有两道保险：

  • tokenizer 没有 chat template 时兜底套 ChatML（src/open_r1/sft.py:94-96，setup_chat_format）；
  • 存盘前把 generation_config.eos_token_id 对齐到 tokenizer 的 EOS（sft.py:131-133），防止训完的模型用 pipeline() 生成时无限续写。

5.3 验收​

这份配方训出的 OpenR1-Distill-7B 对比 DeepSeek 官方蒸馏版（README.md:209-212）：AIME 52.7 vs 51.3、GPQA 52.8 vs 52.4 小幅反超，MATH-500 89.0 vs 93.5、LiveCodeBench 39.4 vs 37.4 互有胜负——蒸馏路线被证实在 7B 量级成立，Step 1 据此宣布完成。

6. 代码地图​

主题文件路径符号名
数据生成管线src/open_r1/generate.pybuild_distilabel_pipeline
生成集群编排slurm/generate.slurm—（bash；Ray job 提交在 176/218 行）
数据混配配置与校验src/open_r1/configs.pyScriptArguments、DatasetMixtureConfig
混配执行src/open_r1/utils/data.pyget_dataset
SFT 训练主线src/open_r1/sft.pymain
蒸馏配方本体recipes/OpenR1-Distill-7B/sft/config_distill.yaml—（YAML）
tokenizer / 模型加载src/open_r1/utils/model_utils.pyget_tokenizer、get_model