您现在的位置是：首页 > 人工智能AI > 正文

人工智能AI

DeepSeekR1Zero中文复现教程来了！

1987web2025-03-01人工智能AI12

Datawhale干货

作者：骆秀韬，Datawhale成员

项目代码可见：unlock-deepseek/Datawhale-R1（https://github/datawhalechina/unlock-deepseek），欢迎关注和 star！

其余所有开源内容见文末。

各位同学好，我是来自 Unlock-DeepSeek 开源项目团队的骆师傅。先说结论，我们（Datawhale X 似然实验室）使用 3 张 A800(80G) 计算卡，花了 20 小时训练时间，做出了可能是国内首批 DeepSeek R1 Zero 的中文复现版本，我们把它叫做 Datawhale-R1，用于 R1 Zero 复现教学。

按照 5.5 元 ~ 7.0 元每小时的价格计算，3 张 A800 花费最低为 3 * 5.5 * 20 = 330 元，预计花费接近 420 元，而TinyZero（https://github/Jiayi-Pan/TinyZero）项目用了 4 张 A800 训练了 8 小时，预计花费为：224 元，这中间的差异可能是由于硬件性能瓶颈和框架差异带来的（我们用的是 Huggingface TRL，TinyZero 使用的是 veRL）。所以建议大家如果真的要复现，请使用 TinyZero 项目，我们出于教育目的使用 TRL 为大家报告这个结果。

另外，不是所有人都能随时随地调用 3 张 A800 的，我们正在努力减小硬件资源要求，让复现工作尽可能平民化（比如在 4090 上跑）。在这里特别感谢：似然实验室，提供本次复现的计算资源，并与 Datawhale 团队合作贡献了本教程。

回到正题，首先回答一个关键问题：为什么这个方案更贵，而我们却选择了它？答案就是：它更符合教育目的，截止本文发布，大部分同学没有足够的资源来亲手体验复现流程，但是我们希望大家能更清楚的看到，复现 R1 Zero 的过程中都发生了什么，真正对复现原理有个大致把握，就算做云玩家也要学到知识，看完骆师傅做一遍就好像自己也做了一遍。

本方案在mini-r1（https://philschmid.de/mini-deepseek-r1）的基础上改进而来。

环境搭建配置基础工具

首先我们要搭建环境，作为手把手教程以及骆师傅的看家本领，我们会在这部分说得细致些。结合国内的实际情况，我们需要的环境信息如下：

暂时无法支持非 Linux 系统（Windows、MacOS）

CUDA > 12.0 （我们使用的是 CUDA 12.4）

Python 建议版本为 3.12（我们使用 Miniforge 管理虚拟环境）

建议使用 Miniforge / Conda 来安装 Pytorch，我们在南方科技大学的开源镜像源测试，下载速度会比官网 pip 安装快不少，请在下面的网址找到适合你硬件的 2.5.1 版本：https://pytorch.org/get-started/previous-versions/，推荐使用 mamba 安装（安装 Miniforge 后直接将 conda 替换为 mamba）

编译安装 flash-attn

接着重头戏就来了，我们需要编译安装 Flash Attention 包，这步非常消耗 CPU 资源，非常不建议CPU核心少的玩家执行。如果你没有办法在有生之年编译完 Flash Attention，可以在 https://github/Dao-AILab/flash-attention/releases/ 找到与你环境对应的编译好的包。（没对应上的话，改环境反而更快，相信我，编译很慢）

这个步骤倒是很简单，执行下面的命令：

pip install packaging

pip install ninja用于加速编译

编译安装 Flash Attention 包

pip install flash-attn --no-build-isolation

注意！如果你的设备CPU核心多，但是运行内存小于 96 GB，请适当设置 MAX_JOBS 的数量，并替换为下面的命令，参考：https://github/Dao-AILab/flash-attentioninstallation-and-features

MAX_JOBS=4 pip install flash-attn --no-build-isolation

按下回车后，可以泡杯咖啡，打开htop看 CPU 疯狂运作，再重新品读一遍《DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning》（https://arxiv.org/abs/2501.12948）

等待 flash-attn 安装完毕后，我们就可以安装其他涉及到的库了，我们提供了一份requirements.txt在Unlock-DeepSeek（https://github/datawhalechina/unlock-deepseek）项目，核心列表如下：

setuptools<71.0.0

transformers==4.48.1

datasets==3.1.0

accelerate==1.3.0

hf-transfer==0.1.9

deepspeed==0.15.4

trl==0.14.0

vllm==0.7.0

modelscope==1.22.3

swanlab==0.4.6

huggingface-hub==0.28.1

大家也可以在这个地址找到我们所有涉及的 Python 包列表：https://swanlab/@anine09/datawhale-r1/runs/4tp31j1zxbm1fshjsi53b/environment/requirements

下载模型和数据集

接下来我们需要下载数据集和模型，在本次实验中，我们使用的数据集为：Jiayi-Pan/Countdown-Tasks-3to4（https://huggingface.co/datasets/Jiayi-Pan/Countdown-Tasks-3to4），模型为：Qwen/Qwen2.5-3B-Instruct（https://huggingface.co/Qwen/Qwen2.5-3B-Instruct），我们目前不建议用小于 3B 的模型（其他社区多次报告，小于 3B 的模型无法学会推理，经过我们的测试，确实！）

数据集下载方式：

exportHF_ENDPOINT=https://hf-mirror更换为国内镜像源，这个只用执行一次，每次重新打开终端就要重新执行，或者写入 .bashrc

下载数据集，替换整个为你自己的内容

huggingface-cli download --repo-type dataset --resume-download Jiayi-Pan/Countdown-Tasks-3to4 --local-dir <你想要存放的路径，比如：dataset>

模型下载方式，哪个速度快用哪个：

方案一，Huggingface 镜像源

下载模型，替换整个为你自己的内容

huggingface-cli download --resume-download Qwen/Qwen2.5-3B-Instruct --local-dir <你想要存放的路径，比如：models>

方案二，ModelScope 下载

新建model_download.py文件，填入以下内容，替换整个为你自己的内容,保存后使用python model_download.py执行下载。

frommodelscopeimportsnapshot_download

model_dir = snapshot_download(Qwen/Qwen2.5-3B-Instruct, cache_dir=<你想要存放的路径，比如：models>, revision=master)

编写配置文件和训练代码

接下来我们需要准备 3 个文件，我们会在Unlock-DeepSeek（https://github/datawhalechina/unlock-deepseek）项目中提供完整的复现文件，方便同学们直接使用。

第一个是 Accelerate 配置文件，用于分布式训练（三张卡）。新建deepspeed_zero3.yaml填入以下内容并保存（不是 DeepSeek，别看错！）。

compute_environment:LOCAL_MACHINE

debug:false

deepspeed_config:

deepspeed_multinode_launcher:standard

offload_optimizer_device:none

offload_param_device:none

zero3_init_flag:true

zero3_save_16bit_model:true

zero_stage:3

distributed_type:DEEPSPEED

downcast_bf16:no

machine_rank:0

main_training_function:main

mixed_precision:bf16

num_machines:1

num_processes:8我们在这里保持常规默认的 8 卡机器，会在后面的启动命令中覆盖新值

rdzv_backend:static

same_network:true

tpu_env:[]

tpu_use_cluster:false

tpu_use_sudo:false

use_cpu:false

一般来说，这个文件内容不需要修改，如果有定制需求，请不要使用这个文件，运行accelerate config自行设定。

在介绍下一个文件之前，我们强烈建议大家使用Swanlab（https://swanlab/）来可视化追踪实验过程，打开：https://swanlab/login ，登录之后点击图中所示的 Quick Start，或者打开：https://swanlab/space/~/settings ，复制 API Key。

在终端输入swanlab login，直接粘贴（你是看不见东西被粘贴上去的），回车，出现类似如下提示就是登录成功。

第二个是 TRL 配置文件，在这里我们会设定训练的超参数。新建Datawhale-R1.yaml填入以下内容，并根据实际情况修改（阅读注释），并保存。

模型参数

model_name_or_path:<你的模型存放的路径，比如：models/Qwen/Qwen2.5-3B-Instruct>

model_revision:main

torch_dtype:bfloat16

attn_implementation:flash_attention_2

bf16:true

tf32:true

output_dir:<你想要模型输出的路径，比如output/Datawhale-R1>

数据集参数

dataset_id_or_path:<你的数据集存放的路径，比如：dataset>

Swanlab 训练流程记录参数

swanlab:true是否开启 Swanlab

workspace:<用户名>

project:<项目名，整个复现项目的名称，例如：Datawhale-R1-by_xxx>

experiment_name:<实验名，某次超参数运行的自定义名称，例如：qwen2.5-3B-lr:5e-7_beta:0.001>

训练参数

max_steps:450最大训练步长

per_device_train_batch_size:1

gradient_accumulation_steps:8

gradient_checkpointing:true

gradient_checkpointing_kwargs:

use_reentrant:false

learning_rate:5.0e-7学习率，调整过，参见下文介绍

lr_scheduler_type:cosine学习率衰减方案

warmup_ratio:0.03学习率预热比率（对于整个步长），好用！

seed:2025随机种子，方便实验复现

GRPO 算法参数

beta:0.001KL 惩罚因子，调整过，参见下文介绍

max_prompt_length:256输入 prompt 最大长度，本实验基本不会有太大变化

max_completion_length:4096输出回答长度，包含推理思维链，设为 4K 比较合适

num_generations:8

use_vllm:true启用 vllm 来加速推理

vllm_device:<计算卡编号，例如：cuda:2>留出一张卡来启用 vllm 推理，参见下文介绍

vllm_gpu_memory_utilization:0.5

Logging arguments

logging_strategy:steps

logging_steps:1

save_strategy:"steps"

save_steps:50每隔多少步保存一次

我们并没有介绍全部参数，如果需要调整，请查阅 Huggingface 相关文档。当然，直接询问 DeepSeek 可能是更快的方式。

这份配置文件中有一些值得大家注意的地方：

learning_rate和beta在 GRPO 的原始论文《DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models》（https://arxiv.org/abs/2402.03300）里分别为1e-6和0.04。在这里我们根据《Unraveling RLHF and Its Variants: Progress and Practical Engineering Insights》（https://hijkzzz.notion.site/unraveling-rlhf-and-its-variants-engineering-insights）将其调整为5e-7和0.001。

vllm_device本实验需要留出一张卡作为 vllm 的推理卡，假设我们手上有 3 张卡（编号cuda: 0, cuda: 1, cuda: 2)，我们需要指定其中一张卡为 vllm 推理卡，例如我们指定最后一张cuda:2。另外，如果你使用了CUDA_VISIBLE_DEVICES情况会有些不一样，比如我们有 8 张卡（编号 cuda:0-7），指定编号为 1、2、3 的卡可见（CUDA_VISIBLE_DEVICES=1,2,3)，这时我们想指定最后一张卡为 vllm 推理卡，则是需要设置为cuda:2，因为设置完可见性后，cuda:1 -> cuda:0，cuda:2 -> cuda:1，cuda:3 -> cuda:2，所以原先的 3 号卡变为了新编号的 2 号卡。

save_steps在mini-r1（https://philschmid.de/mini-deepseek-r1）中是被设为25，但是跑完整个训练后，保存的文件大小达到了 700+ GB！因为不仅包含了模型，还包含了其他卡的优化器状态和其他检查点信息，我们在这里改为50，但仍然要提醒同学们设置成合适自己的大小（训练代码中已经包含结束后保存模型的代码）。

最后，就是创建训练代码文件train_Datawhale-R1.py并保存，我们几乎给每个关键步骤都添加了注释（建议大家从后往前读），在后文我们会再梳理一遍核心步骤。

importlogging

importos

importrandom

importre

fromdataclassesimportdataclass

fromdatetimeimportdatetime

fromtypingimportList

fromdatasetsimportload_dataset

fromswanlab.integration.transformersimportSwanLabCallback

fromtransformersimportAutoTokenizer

fromtransformers.trainer_utilsimportget_last_checkpoint

fromtrlimportGRPOConfig, GRPOTrainer, ModelConfig, TrlParser

@dataclass

classDatasetArguments:

"""数据集参数的数据类"""

数据集 ID 或路径

dataset_id_or_path:str="Jiayi-Pan/Countdown-Tasks-3to4"

数据集拆分

dataset_splits:str="train"

分词器名称或路径

tokenizer_name_or_path:str=None

@dataclass

classSwanlabArguments:

"""SwanLab参数的数据类"""

是否使用 SwanLab

swanlab:bool

SwanLab 用户名

workspace:str

SwanLab 的项目名

project:str

SwanLab 的实验名

experiment_name:str

配置日志记录器

logging.basicConfig(level=logging.INFO)

logger = logging.getLogger(__name__)

logger.setLevel(logging.INFO)

handler = logging.StreamHandler

handler.setFormatter(

logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")

)设置日志格式

logger.addHandler(handler)

defformat_reward_func(completions, **kwargs):

"""

格式奖励函数，检查模型输出格式是否匹配: ......

参数:

completions (list[str]): 生成的输出

list[float]: 奖励分数

"""

初始化奖励列表

rewards = []

遍历生成的输出

forcompletionincompletions:

try:

在生成的输出前添加标签，便于后续正则表达式匹配

completion =""+ completion

ifrandom.random <0.1:1% 的概率将生成输出写入文件

创建生成输出目录（如果不存在）

os.makedirs("completion_samples", exist_ok=True)

log_file = os.path.join("completion_samples","completion_samples.txt")

withopen(log_file,"a")asf:

f.write(f"\n\n==============\n")

f.write(completion)写入生成的输出

定义正则表达式模式，用于匹配和标签

regex =r"^([^<]*(?:<(?!/?think>)[^<]*)*)<\/think>\n([\s\S]*?)<\/answer>$"

match= re.search(regex, completion, re.DOTALL)使用正则表达式进行匹配

ifmatchisNoneorlen(match.groups) !=2:

rewards.append(0.0)如果格式不正确，奖励为 0

else:

rewards.append(1.0)如果格式正确，奖励为 1

exceptException:

rewards.append(0.0)如果发生异常，奖励为 0

returnrewards

defequation_reward_func(completions, target, nums, **kwargs):

"""

方程奖励函数，检查计算结果是否正确，数字是否符合使用要求（每个数字只用一次，只使用所提供的数字）

参数:

completions (list[str]): 生成的输出

target (list[str]): 预期的答案

nums (list[str]): 可用的数字

list[float]: 奖励分数

"""

初始化奖励列表

rewards = []

遍历生成的输出、预期的答案和可用的数字

forcompletion, gt, numbersinzip(completions, target, nums):

try:

在生成的输出前添加标签，便于后续正则表达式匹配

completion =""+ completion

定义正则表达式模式，用于匹配标签

match= re.search(r"(.*?)<\/answer>", completion)

ifmatchisNone:

rewards.append(0.0)如果没有匹配到标签，奖励为 0

continue

equation =match.group(1).strip提取标签中的内容

提取方程中的所有数字

used_numbers = [int(n)forninre.findall(r"\d+", equation)]

检查所有数字是否被使用且只使用一次

ifsorted(used_numbers) !=sorted(numbers):

rewards.append(0.0)

continue

定义允许的字符模式，只允许数字、运算符、括号和空白字符

allowed_pattern =r"^[\d+\-*/.\s]+$"

ifnotre.match(allowed_pattern, equation):

rewards.append(0.0)如果方程包含不允许的字符，奖励为 0

continue

计算方程的结果

result =eval(equation, {"__builtins__":None}, {})

检查方程是否正确且与预期答案匹配（误差小于 1e-5）

ifabs(float(result) -float(gt)) <1e-5:

rewards.append(1.0)如果正确，奖励为 1

10% 的概率将成功的样本写入文件

ifrandom.random <0.10:

创建生成输出目录（如果不存在）

os.makedirs("completion_samples", exist_ok=True)

log_file = os.path.join(

"completion_samples","success_completion_samples.txt"

)

withopen(log_file,"a")asf:

f.write(f"\n\n==============\n")

f.write(completion)写入生成的输出

else:

rewards.append(0.0)如果不正确，奖励为 0

exceptException:

rewards.append(0.0)如果评估失败，奖励为 0

returnrewards

defthought_len_reward_func(completions, **kwargs):

"""

思考长度奖励函数，检查标签的长度是否大于 1000

参数:

completions (list[str]): 生成的输出

list[float]: 奖励分数

"""

初始化奖励列表

rewards = []

遍历生成的输出

forcompletionincompletions:

try:

在生成的输出前添加标签，便于后续正则表达式匹配

completion =""+ completion

定义正则表达式模式，用于匹配标签

match= re.search(r"(.*?)", completion)

如果匹配到标签

ifmatch:

thought_process =match.group(1).strip提取标签中的内容

thought_length =len(thought_process)计算思考过程的长度

ifthought_length >1000:

rewards.append(1.0)如果思考过程长度大于 1000，奖励为 1

else:

rewards.append(0.0)否则奖励为 0

else:

rewards.append(0.0)如果没有匹配到标签，奖励为 0

continue

exceptException:

rewards.append(0.0)如果发生异常，奖励为 0

returnrewards

defget_checkpoint(training_args: GRPOConfig):

"""

获取最后一个检查点

参数:

training_args (GRPOConfig): 训练参数

str: 最后一个检查点的路径，如果没有检查点，则返回 None

"""

last_checkpoint =None

ifos.path.isdir(training_args.output_dir):如果输出目录存在

获取最后一个检查点

last_checkpoint = get_last_checkpoint(training_args.output_dir)

returnlast_checkpoint

定义 GRPO 训练函数

defgrpo_function(

model_args: ModelConfig,

dataset_args: DatasetArguments,

training_args: GRPOConfig,

callbacks:List,

记录模型参数

logger.info(f"Model parameters{model_args}")

记录训练/评估参数

logger.info(f"Training/evaluation parameters{training_args}")

加载分词器

tokenizer = AutoTokenizer.from_pretrained(

(

如果有指定分词器，则使用指定的分词器，否则使用模型名称

dataset_args.tokenizer_name_or_path

ifdataset_args.tokenizer_name_or_path

elsemodel_args.model_name_or_path

revision=model_args.model_revision,使用指定的模型版本

trust_remote_code=model_args.trust_remote_code,允许使用远程代码

)

如果分词器没有填充标记，则使用结束标记作为填充标记

iftokenizer.pad_tokenisNone:

tokenizer.pad_token = tokenizer.eos_token

加载数据集

dataset = load_dataset(

dataset_args.dataset_id_or_path, split=dataset_args.dataset_splits

)

随机选择 50K 个样本，看你喜好定数字，但是数据集有 409K 个样本

dataset = dataset.shuffle(seed=training_args.seed).select(range(50000))

defgenerate_r1_prompt(numbers, target):

"""

生成 R1 Countdown 游戏提示词

参数:

numbers (list[int]): 数字列表

target (int): 目标值

dict: 生成的一个数据样本

"""

定义提示词前缀

r1_prefix = [

{

"role":"user",

"content":f"使用给定的数字{numbers}，创建一个等于{target}的方程。你可以使用基本算术运算（+、-、*、/）一次或多次，但每个数字只能使用一次。在标签中展示你的思考过程，并在标签中返回最终方程，例如 (1 + 2) / 3 。在标签中逐步思考。",

{

"role":"assistant",

"content":"让我们逐步解决这个问题。\n",结尾使用 `` 促使模型开始思考

]

return{

"prompt": tokenizer.apply_chat_template(

r1_prefix, tokenize=False, continue_final_message=True

),提示词，continue_final_message=True 表示将提示词中的最后一个消息继续到最终的输出中

"target": target,

"nums": numbers,

}

将数据集转换为 R1 Countdown 游戏提示词

dataset = dataset.map(lambdax: generate_r1_prompt(x["nums"], x["target"]))

将数据集拆分为训练集和测试集，拆分比例为 9:1

train_test_split = dataset.train_test_split(test_size=0.1)

train_dataset = train_test_split["train"]获取训练集

test_dataset = train_test_split["test"]获取测试集

设置 GRPOTrainer

trainer = GRPOTrainer(

model=model_args.model_name_or_path,模型名称或路径

奖励函数列表，用于计算奖励分数

reward_funcs=[

format_reward_func,格式奖励函数

equation_reward_func,方程奖励函数

thought_len_reward_func,思考长度奖励函数

args=training_args,

train_dataset=train_dataset,

eval_dataset=test_dataset,

callbacks=callbacks,

)

last_checkpoint = get_checkpoint(training_args)检查最后一个检查点

如果检测到检查点且指定从检查点恢复训练，则记录信息

iflast_checkpointisnotNoneandtraining_args.resume_from_checkpointisNone:

logger.info(f"Checkpoint detected, resuming training at{last_checkpoint}.")

logger.info(

f*** Starting training{datetime.now.strftime("%Y-%m-%d %H:%M:%S")}for{training_args.num_train_epochs}epochs***

)

训练模型

train_result = trainer.train(resume_from_checkpoint=last_checkpoint)

记录和保存指标

metrics = train_result.metrics

metrics["train_samples"] =len(train_dataset)

trainer.log_metrics("train", metrics)

trainer.save_metrics("train", metrics)

trainer.save_state

logger.info("*** Training complete ***")

保存模型和分词器

logger.info("*** Save model ***")

trainer.model.config.use_cache =True

trainer.save_model(training_args.output_dir)

logger.info(f"Model saved to{training_args.output_dir}")

training_args.distributed_state.wait_for_everyone等待所有进程加载

tokenizer.save_pretrained(training_args.output_dir)

logger.info(f"Tokenizer saved to{training_args.output_dir}")

logger.info("*** Training complete! ***")

defmain:

"""主函数，用于执行主训练循环"""

解析命令行参数和配置文件

parser = TrlParser((ModelConfig, DatasetArguments, GRPOConfig, SwanlabArguments))

model_args, dataset_args, training_args, swanlab_args = (

parser.parse_args_and_config

)

如果使用 SwanLab，则创建 SwanLab 回调对象，用于训练信息记录

ifswanlab_args.swanlab:

swanlab_callback = SwanLabCallback(

workspace=swanlab_args.workspace,

project=swanlab_args.project,

experiment_name=swanlab_args.experiment_name,

)

callbacks = [swanlab_callback]

else:

callbacks =None

运行主训练循环

grpo_function(model_args, dataset_args, training_args, callbacks=callbacks)

if__name__ =="__main__":

main

启动训练

肯定有一些同学已经等不及要开始跑模型训练了，那启动训练的命令很简单，在终端运行下面的内容（根据自己需求修改），也可以把它保存为train_Datawhale-R1.sh然后在终端运行bash train_Datawhale-R1.sh。

如果你要限制计算卡编号，请在这里设置，例如只使用 cuda:1-3，如果不用限制，就删除下面这行

exportCUDA_VISIBLE_DEVICES=1,2,3

accelerate launch \

--num_processes 2 \

--config_file deepspeed_zero3.yaml \

train_Datawhale-R1.py \

--config Datawhale-R1.yaml

注意：--num_processes是由你希望使用的计算卡数量决定，我们之前在配置文件那里说过，要留一张卡作为 vllm 的推理卡，那么--num_processes的数值应该是你要使用的计算卡数量n-1，例如我有 3 张卡，我的--num_processes应该为2。这里的--num_processes的数值也会把deepspeed_zero3.yaml的num_processes设置的8给覆盖掉。

另外，同样像上文所说，如果你有定制的硬件配置需求，请不要使用--config_file参数。

出现这样的提示就说明模型已经训练起来啦！可以在 Swanlab 看炫酷的训练数据了（手机也能看，特别适合天选炼丹人）。

训练流程详解流程总览

我们来梳理一遍 Datawhale-R1 训练流程：

将提示词输入到 Qwen 2.5 模型。

Qwen 2.5 输出多个带思考的回答（本实验设置为 8，由num_generations参数决定）。

模型的回答分别传入三个奖励函数计算，计算的结果相加。

将奖励值传入 GRPO 策略中，GRPO 根据奖励值来决定如何调整 Qwen 2.5 模型。

重复上述流程（本实验重复了 450 次，由max_steps参数决定）。

有些同学可能不太熟悉强化学习，我们会在后续其他的文章中介绍强化学习相关的概念。在这里我们用一个例子来比喻一下：我们现在假设有一所学校，里面有一个数学老师（GRPO 策略），还有一个班级（Qwen 2.5 模型，我们假设班级中的所有同学能力相同），学校每个月要月考（多步），每次月考是班级根据试卷（提示词，一份试卷只有一道题）写出多份答卷（班级有多个同学，所以会有多份答卷，对应多个带思考的模型回答，这些回答不一定是相同的），这时候数学老师就要去批改这些答卷（奖励函数计算），评卷规则是：

检查答题格式是否规范（格式奖励函数）

解题结果是否正确（方程奖励函数）

解题步骤是否详细（思考长度奖励函数）

最后，把每部分的分数相加，得到多个试卷分数（多个奖励值，用 Python 列表表示，每个回答都对应一个奖励值），数学老师根据班级的月考分数来判断下一步如何调整教学计划（调整模型），来让这个班级在下一次月考中尽可能得到更高的分数。

如果我们说得更细致一点，其实是数学老师会教整个班级看到什么之后写什么，比如看到题目就要写解：，看到x+1=2就要写解得：x=1，力求让组成回答的每一个字都是最合适的（位置要合适，用词也要合适）从而去获得最高的分数。

这里的思考长度奖励函数是我们新加入的，用于鼓励模型进行更长的思考。所以我们应该有个朴素的感受，随着训练的不断进行，Datawhale-R1 的输出格式应该会越来越规范，正确率也会不断提高，思考的长度也会增加。

核心代码介绍

我们稍微介绍一下代码中每个核心步骤的输入输出样例，让大家心里有个底。首先是各种xxx_args参数，它其实就是根据下面这行代码，去获取我们传入的Datawhale-R1.yaml里面的参数。

parser = TrlParser((ModelConfig, DatasetArguments, GRPOConfig, SwanlabArguments))

model_args, dataset_args, training_args, swanlab_args = (

parser.parse_args_and_config

)

你可以看到我们定义了一个SwanlabArguments类，TrlParser会去寻找Datawhale-R1.yaml中跟SwanlabArguments有关的参数，并把它赋值给swanlab_args，由于每个参数名被要求是唯一的，不能重复，所以TrlParser能把不同的参数正确赋值给对应变量（根据ModelConfig, DatasetArguments, GRPOConfig, SwanlabArguments的顺序，赋值给model_args, dataset_args, training_args, swanlab_args）

train_Datawhale-R1.py

@dataclass

classSwanlabArguments:

"""SwanLab参数的数据类"""

是否使用 SwanLab

swanlab:bool

SwanLab 用户名

workspace:str

SwanLab 的项目名

project:str

SwanLab 的实验名

experiment_name:str

Datawhale-R1.yaml

Swanlab 训练流程记录参数

swanlab:true是否开启 Swanlab

workspace:<用户名>

project:<项目名，整个复现项目的名称，例如：Datawhale-R1-by_xxx>

experiment_name:<实验名，某次超参数运行的自定义名称，例如：qwen2.5-3B-lr:5e-7_beta:0.001>

接下来就到了grpo_function里，我们首先来看看我们的数据集长什么样子，我们的任务其实很简单，它很像 24 点游戏，给定若干个数字nums，例如[44, 19, 35]，模型要用四则运算，告诉我们一个方程，它的计算结果正好是target，例如98，详细要求我们在 prompt 中给大家展示。

然后我们的 prompt 如下，利用 Python 的 f-strings 功能来填入具体数值，并且在assistant的结尾加入了\n，来促使我们的模型开始按要求逐步思考。提示词是用 DeepSeek 翻译的 mini-r1 的提示词，咱们中国人阅读中文的速度更快些。

r1_prefix = [

{

"role":"user",

{

"role":"assistant",

"content":"让我们逐步解决这个问题。\n",结尾使用 `` 促使模型开始思考

]

在这里我们会把 prompt 转换为 Qwen 2.5 的提示词模版，让它以更熟悉的方式来接收提示词，并且我们把让我们逐步解决这个问题。\n作为模型输出的开头，让它接着续写。用 Python 字典的方式返回样本，这样 TRL 会在调用奖励函数的时候，帮我们把键名设为为对应的参数；另外，TRL 会把模型的多个输出设为completions。

return{

"prompt": tokenizer.apply_chat_template(

r1_prefix, tokenize=False, continue_final_message=True

),提示词，continue_final_message=True 表示将提示词中的最后一个消息继续到最终的输出中

"target": target,

"nums": numbers,

}

map方法会帮我们把实际的nums和target填入到 prompt 里，我们根据上面举的例子，来看一个具体的提示词：

将数据集转换为 R1 Countdown 游戏提示词

dataset = dataset.map(lambdax: generate_r1_prompt(x["nums"], x["target"]))

举例

nums = [44,19,35]

target =98

r1_prefix = {

"role":"user",

"content":f"使用给定的数字 [44, 19, 35]，创建一个等于 98 的方程。你可以使用基本算术运算（+、-、*、/）一次或多次，但每个数字只能使用一次。在标签中展示你的思考过程，并在标签中返回最终方程，例如 (1 + 2) / 3 。在标签中逐步思考。",

{

"role":"assistant",

"content":"让我们逐步解决这个问题。\n",结尾使用 `` 促使模型开始思考

转换为 Qwen 提示词模版后

我们最后来看一个奖励函数的例子，TRL 将多个模型输出变成一个列表，叫做completions，并将数据集中的其他内容根据键名传入到对应参数。所以我们需要使用 for 循环遍历所有的completions，并对每个输出进行判断打分，最后返回每个输出的得分列表reward给GRPO 策略（例如：[0.0, 1.0, 0.0]），让其判断下一步如何调整。

defequation_reward_func(completions, target, nums, **kwargs):

"""

参数:

completions (list[str]): 生成的输出

target (list[str]): 预期的答案

nums (list[str]): 可用的数字

list[float]: 奖励分数

"""

初始化奖励列表

rewards = []

遍历生成的输出、预期的答案和可用的数字

forcompletion, gt, numbersinzip(completions, target, nums):

...进行一些 rewards.append 操作

returnrewards

OK，我们对复现流程的介绍就大致结束了，我们会在文末提供完整的文档，开源我们的复现工作。

训练结果解读

现在我们来看看模型表现出了什么有意思的现象，提前声明，这不是严谨的科学研究，会有很多分析漏洞。首先我们使用了学习率预热和学习率衰减，在训练前期学习率都很大，后期慢慢衰减下来。

对比下面两张图，我们发现模型前期学习输出格式的速度很快，大概 20 到 30 步就能学得很好。但是后来由于我们的思考长度奖励函数，模型的输出长度被拉长，发生严重的重复现象，导致超出 4096 的输出被截断，格式不完整，格式奖励函数的奖励值就大幅下降，后面模型又开始缩短输出，稳定在 300 到 400，又恢复到正确格式。

模型的不断重复输出看着其实挺可怕的，Visual Studio Code 会匹配相同字符并高亮，大家可以看看，右侧红框的缩略图几乎都是重复的回答。

我们发现，模型被鼓励拉长输出的时候，计算正确率也在提升，所以我们有个不严谨的判断，似乎拉长模型输出，能带一定的计算正确率的提升。观察下图可以发现，在 120 步时，模型的输出在越变越长，平均输出长度已经被拉到 400 左右，越来越多的输出已经超过 1000，方程计算正确率也在逐步升高，但是这时已经发生一些重复问题导致格式错误。

其实从上图我们也可以看到 GRPO 已经意识到重复问题带来的奖励值下降，它在 200 步左右开始逐步限制模型输出长度，而这时模型的计算正确率也保持在 0.3 到 0.4 左右。

我们还发现，在训练初期，你会看到比较明显的方程奖励提升，而输出长度不断减小。模型似乎有一种趋向于缩短思考长度的趋势，所以我们引入思考长度奖励函数来对抗这种趋势，我们把它解释为模型计算能力提升之后，就像学霸一眼秒杀题目一样，模型不想输出更多废话来解释解题过程。

在训练开始 1 分钟左右，我们就观察到下面的输出，还以为我们重现了 Aha Moment。后来证明其实不是，Qwen 2.5 很喜欢反复试错、验算，反复试错很容易导致上文提及的重复输出问题。

我们发现了另一种语言混用现象，哈哈哈。current n. 电流; adj. 当前的。

所以结论就是，我们没有复现 Aha Moment。其实在观察大量 Qwen 2.5 的输出之后，一种直觉告诉我，可能 Aha Moment 跟模型本身的输出风格相关，网友都说 DeepSeek 文风很锐利、很活泼，但是 Qwen 2.5 给人的感觉总是冷静、平和。简单做了一个不严谨的测试，可能能够佐证这个想法，我要求两个模型用 Aha Moment 的语气跟我说话，再随便回复了一个字，观察两个模型本身对 Aha Moment 的映射是会输出什么。我们另外测试了 Llama 和 MiniCPM，它们的输出风格都跟 Qwen 很接近，试图像说教一样给你做比喻，所以我大胆判断，可能写武侠小说的大模型更容易观察到 Aha Moment。