DeepSeekR1Zero中文复现教程来了

干货

作者 : 骆秀韬,成员

项目代码可见:unlock-deepseek/-R1(https://github.com/datawhalechina/unlock-deepseek),欢迎关注和 star!

其余所有开源内容见文末。本文实现细节有bug,已在文末网址更正。

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

DeepSeekR1Zero中文复现教程来了-2

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

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

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

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

环境搭建

配置基础工具

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

暂时无法支持非 Linux 系统(Windows、MacOS)

CUDA > 12.0 (我们使用的是 CUDA 12.4)

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

Pytorch 版本为 2.5.1 (GPU版本,请使用 torch.cuda.is_available() 检查能否正常识别 GPU 设备)

建议使用 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.com/Dao-AILab/flash-attention/releases/ 找到与你环境对应的编译好的包。(没对应上的话,改环境反而更快,相信我,编译很慢)

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

pip install packagingpip install ninja # 用于加速编译# 编译安装 Flash Attention 包pip install flash-attn --no-build-isolation# 注意!如果你的设备CPU核心多,但是运行内存小于 96 GB,请适当设置 MAX_JOBS 的数量,并替换为下面的命令,参考:https://github.com/Dao-AILab/flash-attention#installation-and-featuresMAX_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.com/datawhalechina/unlock-deepseek) 项目,核心列表如下:

setuptools<71.0.0transformers==4.48.1datasets==3.1.0accelerate==1.3.0hf-transfer==0.1.9deepspeed==0.15.4trl==0.14.0vllm==0.7.0modelscope==1.22.3swanlab==0.4.6huggingface-hub==0.28.1

大家也可以在这个地址找到我们所有涉及的 Python 包列表:https://swanlab.cn/@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 的模型无法学会推理,经过我们的测试,确实!)

数据集下载方式:

export HF_ENDPOINT=https://hf-mirror.com # 更换为国内镜像源,这个只用执行一次,每次重新打开终端就要重新执行,或者写入 .bashrc# 下载数据集,替换整个 <xxx> 为你自己的内容huggingface-cli download --repo-type dataset --resume-download Jiayi-Pan/Countdown-Tasks-3to4 --local-dir <你想要存放的路径,比如:dataset>

模型下载方式,哪个速度快用哪个:

方案一,Huggingface 镜像源

# 下载模型,替换整个 <xxx> 为你自己的内容huggingface-cli download --resume-download Qwen/Qwen2.5-3B-Instruct --local-dir <你想要存放的路径,比如:models>

方案二,ModelScope 下载

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

from modelscope import snapshot_downloadmodel_dir = snapshot_download('Qwen/Qwen2.5-3B-Instruct', cache_dir='<你想要存放的路径,比如:models>', revision='master')

编写配置文件和训练代码

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

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

compute_environment: LOCAL_MACHINEdebug: falsedeepspeed_config:  deepspeed_multinode_launcher: standard  offload_optimizer_device: none  offload_param_device: none  zero3_init_flag: true  zero3_save_16bit_model: true  zero_stage: 3distributed_type: DEEPSPEEDdowncast_bf16: 'no'machine_rank: 0main_training_function: mainmixed_precision: bf16num_machines: 1num_processes: 8 # 我们在这里保持常规默认的 8 卡机器,会在后面的启动命令中覆盖新值rdzv_backend: staticsame_network: truetpu_env: []tpu_use_cluster: falsetpu_use_sudo: falseuse_cpu: false

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

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

DeepSeekR1Zero中文复现教程来了-3

DeepSeekR1Zero中文复现教程来了-4

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

DeepSeekR1Zero中文复现教程来了-5

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

# 模型参数model_name_or_path: <你的模型存放的路径,比如:models/Qwen/Qwen2.5-3B-Instruct>model_revision: maintorch_dtype: bfloat16attn_implementation: flash_attention_2bf16: truetf32: trueoutput_dir: <你想要模型输出的路径,比如 output/-R1># 数据集参数dataset_id_or_path: <你的数据集存放的路径,比如:dataset># Swanlab 训练流程记录参数swanlab: true # 是否开启 Swanlab workspace: <用户名>project: <项目名,整个复现项目的名称,例如:-R1-by_xxx>experiment_name: <实验名,某次超参数运行的自定义名称,例如:qwen2.5-3B-lr:5e-7_beta:0.001># 训练参数max_steps: 450 # 最大训练步长per_device_train_batch_size: 1gradient_accumulation_steps: 8gradient_checkpointing: truegradient_checkpointing_kwargs:  use_reentrant: falselearning_rate: 5.0e-7 # 学习率,调整过,参见下文介绍lr_scheduler_type: cosine # 学习率衰减方案warmup_ratio: 0.03 # 学习率预热比率(对于整个步长),好用!seed: 2025 # 随机种子,方便实验复现# GRPO 算法参数beta: 0.001 # KL 惩罚因子,调整过,参见下文介绍max_prompt_length: 256 # 输入 prompt 最大长度,本实验基本不会有太大变化max_completion_length: 4096 # 输出回答长度,包含推理思维链,设为 4K 比较合适num_generations: 8use_vllm: true # 启用 vllm 来加速推理vllm_device: <计算卡编号,例如:cuda:2> # 留出一张卡来启用 vllm 推理,参见下文介绍vllm_gpu_memory_utilization: 0.5# Logging argumentslogging_strategy: stepslogging_steps: 1save_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://www.philschmid.de/mini-deepseek-r1) 中是被设为 25 ,但是跑完整个训练后,保存的文件大小达到了 700+ GB!因为不仅包含了模型,还包含了其他卡的优化器状态和其他检查点信息,我们在这里改为 50 ,但仍然要提醒同学们设置成合适自己的大小(训练代码中已经包含结束后保存模型的代码)。

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

import loggingimport osimport randomimport refrom dataclasses import dataclassfrom datetime import datetimefrom typing import Listfrom datasets import load_datasetfrom swanlab.integration.transformers import SwanLabCallbackfrom transformers import AutoTokenizerfrom transformers.trainer_utils import get_last_checkpointfrom trl import GRPOConfig, GRPOTrainer, ModelConfig, TrlParser@dataclassclass DatasetArguments:    """数据集参数的数据类"""    # 数据集 ID 或路径    dataset_id_or_path: str = "Jiayi-Pan/Countdown-Tasks-3to4"    # 数据集拆分    dataset_splits: str = "train"    # 分词器名称或路径    tokenizer_name_or_path: str = None@dataclassclass SwanlabArguments:    """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)def format_reward_func(completions, **kwargs):    """    格式奖励函数,检查模型输出格式是否匹配: <think>...</think><answer>...</answer>    参数:        completions (list[str]): 生成的输出    返回:        list[float]: 奖励分数    """    # 初始化奖励列表    rewards = []    # 遍历生成的输出    for completion in completions:        try:            # 在生成的输出前添加<think>标签,便于后续正则表达式匹配            completion = "<think>" + completion            if random.random() < 0.1:  # 1% 的概率将生成输出写入文件                # 创建生成输出目录(如果不存在)                os.makedirs("completion_samples", exist_ok=True)                log_file = os.path.join("completion_samples", "completion_samples.txt")                with open(log_file, "a") as f:                    f.write(f"nn==============n")                    f.write(completion)  # 写入生成的输出            # 定义正则表达式模式,用于匹配 <think> 和 <answer> 标签            regex = r"^<think>([^<]*(?:<(?!/?think>)[^<]*)*)</think>n<answer>([sS]*?)</answer>$"            match = re.search(regex, completion, re.DOTALL)  # 使用正则表达式进行匹配            if match is None or len(match.groups()) != 2:                rewards.append(0.0)  # 如果格式不正确,奖励为 0            else:                rewards.append(1.0)  # 如果格式正确,奖励为 1        except Exception:            rewards.append(0.0)  # 如果发生异常,奖励为 0    return rewardsdef equation_reward_func(completions, target, nums, **kwargs):    """    方程奖励函数,检查计算结果是否正确,数字是否符合使用要求(每个数字只用一次,只使用所提供的数字)    参数:        completions (list[str]): 生成的输出        target (list[str]): 预期的答案        nums (list[str]): 可用的数字    返回:        list[float]: 奖励分数    """    # 初始化奖励列表    rewards = []    # 遍历生成的输出、预期的答案和可用的数字    for completion, gt, numbers in zip(completions, target, nums):        try:            # 在生成的输出前添加 <think> 标签,便于后续正则表达式匹配            completion = "<think>" + completion            # 定义正则表达式模式,用于匹配 <answer> 标签            match = re.search(r"<answer>(.*?)</answer>", completion)            if match is None:                rewards.append(0.0)  # 如果没有匹配到 <answer> 标签,奖励为 0                continue            equation = match.group(1).strip()  # 提取 <answer> 标签中的内容            # 提取方程中的所有数字            used_numbers = [int(n) for n in re.findall(r"d+", equation)]            # 检查所有数字是否被使用且只使用一次            if sorted(used_numbers) != sorted(numbers):                rewards.append(0.0)                continue            # 定义允许的字符模式,只允许数字、运算符、括号和空白字符            allowed_pattern = r"^[d+-*/().s]+$"            if not re.match(allowed_pattern, equation):                rewards.append(0.0)  # 如果方程包含不允许的字符,奖励为 0                continue            # 计算方程的结果            result = eval(equation, {"__builtins__": None}, {})            # 检查方程是否正确且与预期答案匹配(误差小于 1e-5)            if abs(float(result) - float(gt)) < 1e-5:                rewards.append(1.0)  # 如果正确,奖励为 1                # 10% 的概率将成功的样本写入文件                if random.random() < 0.10:                    # 创建生成输出目录(如果不存在)                    os.makedirs("completion_samples", exist_ok=True)                    log_file = os.path.join(                        "completion_samples", "success_completion_samples.txt"                    )                    with open(log_file, "a") as f:                        f.write(f"nn==============n")                        f.write(completion)  # 写入生成的输出            else:                rewards.append(0.0)  # 如果不正确,奖励为 0        except Exception:            rewards.append(0.0)  # 如果评估失败,奖励为 0    return rewardsdef thought_len_reward_func(completions, **kwargs):    """    思考长度奖励函数,检查 <think> 标签的长度是否大于 1000    参数:        completions (list[str]): 生成的输出    返回:        list[float]: 奖励分数    """    # 初始化奖励列表    rewards = []    # 遍历生成的输出    for completion in completions:        try:            # 在生成的输出前添加 <think> 标签,便于后续正则表达式匹配            completion = "<think>" + completion            # 定义正则表达式模式,用于匹配 <think> 标签            match = re.search(r"<think>(.*?)</think>", completion)            # 如果匹配到 <think> 标签            if match:                thought_process = match.group(1).strip()  # 提取 <think> 标签中的内容                thought_length = len(thought_process)  # 计算思考过程的长度                if thought_length > 1000:                    rewards.append(1.0)  # 如果思考过程长度大于 1000,奖励为 1                else:                    rewards.append(0.0)  # 否则奖励为 0            else:                rewards.append(0.0)  # 如果没有匹配到 <think> 标签,奖励为 0                continue        except Exception:            rewards.append(0.0)  # 如果发生异常,奖励为 0    return rewardsdef get_checkpoint(training_args: GRPOConfig):    """    获取最后一个检查点    参数:        training_args (GRPOConfig): 训练参数    返回:        str: 最后一个检查点的路径,如果没有检查点,则返回 None    """    last_checkpoint = None    if os.path.isdir(training_args.output_dir):  # 如果输出目录存在        # 获取最后一个检查点        last_checkpoint = get_last_checkpoint(training_args.output_dir)    return last_checkpoint# 定义 GRPO 训练函数def grpo_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            if dataset_args.tokenizer_name_or_path            else model_args.model_name_or_path        ),        revision=model_args.model_revision,  # 使用指定的模型版本        trust_remote_code=model_args.trust_remote_code,  # 允许使用远程代码    )    # 如果分词器没有填充标记,则使用结束标记作为填充标记    if tokenizer.pad_token is None:        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))    def generate_r1_prompt(numbers, target):        """        生成 R1 Countdown 游戏提示词        参数:            numbers (list[int]): 数字列表            target (int): 目标值        返回:            dict: 生成的一个数据样本        """        # 定义提示词前缀        r1_prefix = [            {                "role": "user",                "content": f"使用给定的数字 {numbers},创建一个等于 {target} 的方程。你可以使用基本算术运算(+、-、*、/)一次或多次,但每个数字只能使用一次。在 <think> </think> 标签中展示你的思考过程,并在 <answer> </answer> 标签中返回最终方程,例如 <answer> (1 + 2) / 3 </answer>。在 <think> 标签中逐步思考。",            },            {                "role": "assistant",                "content": "让我们逐步解决这个问题。n<think>",  # 结尾使用 `<think>` 促使模型开始思考            },        ]        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(lambda x: 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)  # 检查最后一个检查点    # 如果检测到检查点且指定从检查点恢复训练,则记录信息    if last_checkpoint is not None and training_args.resume_from_checkpoint is None:        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! ***")def main():    """主函数,用于执行主训练循环"""    # 解析命令行参数和配置文件    parser = TrlParser((ModelConfig, DatasetArguments, GRPOConfig, SwanlabArguments))    model_args, dataset_args, training_args, swanlab_args = (        parser.parse_args_and_config()    )    # 如果使用 SwanLab,则创建 SwanLab 回调对象,用于训练信息记录    if swanlab_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_-R1.sh 然后在终端运行 bash train_-R1.sh 。

# 如果你要限制计算卡编号,请在这里设置,例如只使用 cuda:1-3,如果不用限制,就删除下面这行export CUDA_VISIBLE_DEVICES=1,2,3accelerate launch     --num_processes 2     --config_file deepspeed_zero3.yaml     train_-R1.py     --config -R1.yaml

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

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

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

DeepSeekR1Zero中文复现教程来了-6

DeepSeekR1Zero中文复现教程来了-7

训练流程详解

流程总览

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

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

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

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

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

重复上述流程(本实验重复了 450 次,由 max_steps 参数决定)。

DeepSeekR1Zero中文复现教程来了-8

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

检查答题格式是否规范(格式奖励函数)

解题结果是否正确(方程奖励函数)

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

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

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

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

核心代码介绍

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

parser = TrlParser((ModelConfig, DatasetArguments, GRPOConfig, SwanlabArguments))model_args, dataset_args, training_args, swanlab_args = (    parser.parse_args_and_config())

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

# train_-R1.py@dataclassclass SwanlabArguments:    """SwanLab参数的数据类"""    # 是否使用 SwanLab    swanlab: bool    # SwanLab 用户名    workspace: str    # SwanLab 的项目名    project: str    # SwanLab 的实验名    experiment_name: str

# -R1.yaml# Swanlab 训练流程记录参数swanlab: true # 是否开启 Swanlab workspace: <用户名>project: <项目名,整个复现项目的名称,例如:-R1-by_xxx>experiment_name: <实验名,某次超参数运行的自定义名称,例如:qwen2.5-3B-lr:5e-7_beta:0.001>

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

DeepSeekR1Zero中文复现教程来了-9

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

r1_prefix = [    {        "role": "user",        "content": f"使用给定的数字 {numbers},创建一个等于 {target} 的方程。你可以使用基本算术运算(+、-、*、/)一次或多次,但每个数字只能使用一次。在 <think> </think> 标签中展示你的思考过程,并在 <answer> </answer> 标签中返回最终方程,例如 <answer> (1 + 2) / 3 </answer>。在 <think> 标签中逐步思考。",    },    {        "role": "assistant",        "content": "让我们逐步解决这个问题。n<think>",  # 结尾使用 `<think>` 促使模型开始思考    },]

在这里我们会把 prompt 转换为 Qwen 2.5 的提示词模版,让它以更熟悉的方式来接收提示词,并且我们把 让我们逐步解决这个问题。n <think> 作为模型输出的开头,让它接着续写。用 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(lambda x: generate_r1_prompt(x["nums"], x["target"]))# 举例nums = [44, 19, 35]target = 98r1_prefix = {    "role": "user",    "content": f"使用给定的数字 [44, 19, 35],创建一个等于 98 的方程。你可以使用基本算术运算(+、-、*、/)一次或多次,但每个数字只能使用一次。在 <think> </think> 标签中展示你的思考过程,并在 <answer> </answer> 标签中返回最终方程,例如 <answer> (1 + 2) / 3 </answer>。在 <think> 标签中逐步思考。",},{    "role": "assistant",    "content": "让我们逐步解决这个问题。n<think>",  # 结尾使用 `<think>` 促使模型开始思考},# 转换为 Qwen 提示词模版后prompt = "<|im_start|>systemnYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>n<|im_start|>usern使用给定的数字 [44, 19, 35],创建一个等于 98 的方程。你可以使用基本算术运算(+、-、*、/)一次或多次,但每个数字只能使用一次。在 <think> </think> 标签中展示你的思考过程,并在 <answer> </answer> 标签中返回最终方程,例如 <answer> (1 + 2) / 3 </answer>。在 <think> 标签中逐步思考。<|im_end|>n<|im_start|>assistantn让我们逐步解决这个问题。n<think>" # 模型将在 n<think> 后续写

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

def equation_reward_func(completions, target, nums, **kwargs):    """    参数:        completions (list[str]): 生成的输出        target (list[str]): 预期的答案        nums (list[str]): 可用的数字    返回:        list[float]: 奖励分数    """    # 初始化奖励列表    rewards = []    # 遍历生成的输出、预期的答案和可用的数字    for completion, gt, numbers in zip(completions, target, nums):        ... # 进行一些 rewards.append() 操作    return rewards

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

训练结果解读

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

DeepSeekR1Zero中文复现教程来了-10

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

DeepSeekR1Zero中文复现教程来了-11

DeepSeekR1Zero中文复现教程来了-12

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

DeepSeekR1Zero中文复现教程来了-13

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

DeepSeekR1Zero中文复现教程来了-14

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

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

DeepSeekR1Zero中文复现教程来了-15

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

DeepSeekR1Zero中文复现教程来了-16

DeepSeekR1Zero中文复现教程来了-17

DeepSeekR1Zero中文复现教程来了-18

DeepSeekR1Zero中文复现教程来了-19

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

DeepSeekR1Zero中文复现教程来了-20

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

DeepSeekR1Zero中文复现教程来了-21

DeepSeekR1Zero中文复现教程来了-22

我们会一同公布模型输出的采样文本文件,大家也可以在里面找到一些我们还没有发现的新奇玩意,欢迎向 Unlock-DeepSeek 团队报告你的发现。

展望

在本文写作前一天,我们发现 另一组团队也公开了他们的 Qwen 2.5 7B 的 R1 Zero 复现结果 (https://zhuanlan.zhihu.com/p/21290410831) ,他们也观察到了很多有趣的结果,虽然他们的曲线非常震荡,但是也稍微能看出一点佐证我们观点的证据:似乎拉长模型输出,能带一定的计算正确率的提升。他们的工作非常棒!我们就不用去验证 7B 模型的性能了,非常环保,节能减排。大家也可以追踪观察社区其他小组的复现报告,相信开源社区的力量!

最后嘱咐一些要点,

Math 模型不太好用,它有固有的数学输出会影响格式奖励,可能需要更长的步长才能纠正,不环保,训了一会我就停了。

DeepSeekR1Zero中文复现教程来了-23

DeepSeekR1Zero中文复现教程来了-24

小于 3B 的模型真不好用,没什么必要再试验了,DeepSeek 官方蒸馏的 1.5B 的推理也很烂,小模型承受了太多它不该承受的东西。我们甚至还在 0.5B 的模型看到了俄语,但是找不到图了。

这种训练方式用来规范模型输出格式特别好用。

Jian Hu 报告 GRPO 有严重震荡问题 (https://zhuanlan.zhihu.com/p/14888098807) ,或许大家可以试试其他算法。

如果你的资源充足,可以试试更大的模型,希望在开源社区能够见到大家的新发现。

TRL 目前的 LoRA 模块有严重 Bug,请不要使用。

最后一点,要复现,请用 TinyZero,省钱!

完整文件获取

Unlock-DeepSeek 团队后续会陆续发布更多关于 DeepSeek 相关工作解读的文章(马上就会发布 GRPO 解读文章),敬请关注,我们下次再见!

Unlock-DeepSeek 项目主页:https://datawhalechina.github.io/unlock-deepseek/

Github 仓库:https://github.com/datawhalechina/unlock-deepseek

Gitee 国内仓库:https://gitee.com/anine09/unlock-deepseek

Swanlab 实验数据:https://swanlab.cn/@anine09/datawhale-r1/overview

模型会在晚些时候上传 HuggingFace 和 ModelScope,并在项目中公布,虽然模型本身没什么用。

复现文件在 -R1 文件夹。

Unlock-DeepSeek 项目目前并不完善,并且正在快速迭代,请持续关注。

THE END
分享
二维码
< <上一篇
下一篇>>