618ZXW

DeepSeek R1 Zero を再現するための中国語チュートリアルがここにあります!

オリジナル:Luo Shifu(Datawhale)

Datawhaleのヒント

著者 Luo Xiutao、Datawhaleメンバー

プロジェクトのコードは、unlock-deepseek/Datawhale-R1 (https://github.com/datawhalechina/unlock-deepseek) にあります。ぜひフォローしてスターを付けてください!

その他のオープンソース コンテンツはすべてこの記事の最後にあります。


皆さん、こんにちは。Unlock-DeepSeekオープンソースプロジェクトチームのLuoです。まとめると、私たち(Datawhale X Likelihood Lab)はA800(80G)コンピューティングカード3枚を使用し、20時間のトレーニングを経て、おそらく中国で初となるDeepSeek R1 Zeroの複製を作成しました。このマシンをDatawhale-R1と名付け、R1 Zeroの複製学習に使用しています。

1時間あたり5.5~7.0元の価格に基づくと、3基のA800 GPUを使用したトレーニングの最小コストは3 * 5.5 * 20 = 330元となり、推定コストは約420元となります。TinyZeroプロジェクト(https://github.com/Jiayi-Pan/TinyZero)は、4基のA800 GPUを8時間使用し、推定コストは224元でした。この差は、ハードウェアパフォーマンスのボトルネックとフレームワークの違い(TinyZeroではveRLを使用、TinyZeroではHuggingface TRLを使用)によるものと考えられます。したがって、この結果を実際に再現したい場合は、TinyZeroプロジェクトの使用をお勧めします。この結果は教育目的でTRLを使用して報告しています。

さらに、誰もが3基のA800 GPUを容易に利用できるわけではありません。私たちは、ハードウェアリソースの要件を削減し、再現プロセスを可能な限り容易にする(例えば、4090で実行できるようにする)よう取り組んでいます。この再現に必要な計算リソースを提供してくださったLikelihood Labと、このチュートリアルの作成にご協力いただいたDatawhaleチームに深く感謝いたします。

本題に戻りましょう。まずは重要な疑問に答えましょう。なぜ私たちはこの高価なソリューションを選んだのでしょうか?答えは、教育目的により合致するからです。この記事の公開時点では、ほとんどの学生は複製プロセスを実際に体験するのに十分なリソースを持っていません。しかし、R1 Zeroの複製中に何が起こるのかをより明確に理解し、複製の原理を真に理解していただければ幸いです。たとえ「クラウドユーザー」であっても、何かを学べるはずです。羅師匠のデモンストレーションを見ていると、まるで自分でやっているかのような感覚になります。

このソリューションは、mini-r1 (https://www.philschmid.de/mini-deepseek-r1) の改良版です。

環境設定

基本ツールを設定する

まず、環境を構築する必要があります。これはステップバイステップのチュートリアルであり、Luo師匠の得意技の一つでもあるため、この部分について詳しく説明します。中国の実情を踏まえると、必要な環境は以下のとおりです。

Linux 以外のシステム (Windows、macOS) は現在サポートされていません。

  • CUDA > 12.0 (CUDA 12.4 を使用しています)
  • 推奨される Python バージョンは 3.12 です (仮想環境の管理には Miniforge を使用します)。
  • PyTorch バージョンは 2.5.1 です (GPU バージョン。GPU デバイスが正しく認識できるかどうかを確認するには、torch.cuda.is_available() を使用してください)。

PyTorchのインストールにはMiniforge/Condaの使用をお勧めします。南方科技大学のオープンソースミラーで行ったテストでは、公式サイトからのpipインストールよりもダウンロード速度が大幅に向上しました。お使いのハードウェアに適したバージョン2.5.1は、以下のURLでご確認ください: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 packaging pip install ninja # 用于加速编译

Flash Attentionパッケージをコンパイルしてインストールする

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

注意!デバイスのCPUコア数は多いものの、RAMが96GB未満の場合は、MAX_JOBSの値を調整し、以下のコマンドに置き換えてください。詳細はhttps://github.com/Dao-AILab/flash-attention#installation-and-featuresをご覧ください。

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

Enter キーを押した後、コーヒーを一杯淹れて、htop を開いて CPU が狂ったように動作するのを観察してから、「DeepSeek-R1: 強化学習による LLM の推論能力のインセンティブ化」(https://arxiv.org/abs/2501.12948) を再度読んでみてください。

flash-attn のインストールが完了したら、その他の関連ライブラリをインストールできます。Unlock-DeepSeek プロジェクト (https://github.com/datawhalechina/unlock-deepseek) では requirements.txt ファイルを提供しており、コアとなるライブラリのリストは以下のとおりです。

 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.cn/@anine09/datawhale-r1/runs/4tp31j1zxbm1fsh...

モデルとデータセットをダウンロードする

次に、データセットとモデルをダウンロードする必要があります。今回の実験では、データセット: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>

モデルを最も速くダウンロードする方法を使用します。

  • オプション1:Huggingfaceミラーソース

モデルをダウンロードし、`<xxx>` 全体を独自のコンテンツに置き換えます。

 huggingface-cli download --resume-download Qwen/Qwen2.5-3B-Instruct --local-dir <你想要存放的路径,比如:models>
  • オプション2: ModelScopeをダウンロードする

model_download.py という名前の新しいファイルを作成し、次の内容を入力し、<xxx> 全体を独自の内容に置き換えて保存し、python model_download.py を使用してダウンロードを実行します。

 from modelscope import snapshot\_download model\_dir = snapshot\_download('Qwen/Qwen2.5-3B-Instruct', cache\_dir='<你想要存放的路径,比如:models>', revision='master')

設定ファイルとトレーニングコードの作成

次に、3つのファイルを準備する必要があります。学生が直接使用できるように、Unlock-DeepSeekプロジェクト(https://github.com/datawhalechina/unlock-deepseek)で完全な再現ファイルを提供します。

  • 1つ目は、分散学習(GPU 3基)に使用する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

通常、このファイルを変更する必要はありません。カスタマイズが必要な場合は、このファイルを使用せず、accelerated config を実行してご自身で設定してください。

次のドキュメントをご紹介する前に、実験プロセスを可視化・追跡するためにSwanlab(https://swanlab.cn/)をご利用いただくことを強くお勧めします。https://swanlab.cn/loginにアクセスしてログインし、画像のように「クイックスタート」をクリックするか、https://swanlab.cn/space/~/settingsにアクセスしてAPIキーをコピーしてください。

ターミナルで「swanlab login」と入力し、テキストを直接貼り付けます(貼り付けられた内容は表示されません)。Enterキーを押します。以下のようなメッセージが表示されれば、ログイン成功です。

  • 2つ目は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.001 # KL 惩罚因子,调整过,参见下文介绍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\_strategy: steps logging\_steps: 1 save\_strategy: "steps" save\_steps: 50 # 每隔多少步保存一次

すべてのパラメータを網羅しているわけではありません。調整が必要な場合は、Huggingfaceのドキュメントを参照してください。もちろん、DeepSeekに直接問い合わせる方が速いかもしれません。

この構成ファイルには注目すべき点がいくつかあります。

  • オリジナルのGRPO論文「DeepSeekMath: オープン言語モデルにおける数学的推論の限界の探求」(https://arxiv.org/abs/2402.03300) では、`learning_rate` と `beta` はそれぞれ 1e-6 と 0.04 です。ここでは、「Unraveling RLHF and Its Variants: Progress and Practical Engineering Insights」(https://hijkzzz.notion.site/unraveling-rlhf-and-its-variants-...) に基づき、それぞれ 5e-7 と 0.001 に調整します。
  • この実験では、1枚のカードをVLLM推論カードとして予約する必要があります。3枚のカード(cuda: 0、cuda: 1、cuda: 2)があると仮定すると、そのうちの1枚を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 になるためです。
  • mini-r1 (https://www.philschmid.de/mini-deepseek-r1) では、`save_steps` は 25 に設定されていますが、トレーニング全体を実行すると、保存されたファイルサイズが 700 GB を超えました。これは、モデルだけでなく、他の GPU のオプティマイザーステータスやその他のチェックポイント情報も含まれているためです。ここでは 50 に変更しましたが、受講者の皆様には、ご自身のニーズに合わせて適切なサイズに設定することをお勧めしています(トレーニングコードには、完了後にモデルを保存するコードが既に含まれています)。
  • 最後に、トレーニングコードファイル train_Datawhale-R1.py を作成して保存しました。ほぼすべての主要なステップにコメントを追加しました(最後から始めまで読むことをお勧めします)。コアとなるステップについては後ほど改めて確認します。
 import logging import os import random import re from dataclasses import dataclass from datetime import datetime from typing import List from datasets import load\_dataset from swanlab.integration.transformers import SwanLabCallback from transformers import AutoTokenizer from transformers.trainer\_utils import get\_last\_checkpoint from trl import GRPOConfig, GRPOTrainer, ModelConfig, TrlParser @dataclass class DatasetArguments:    """数据集参数的数据类"""    # 数据集 ID 或路径    dataset\_id\_or\_path: str = "Jiayi-Pan/Countdown-Tasks-3to4"    # 数据集拆分    dataset\_splits: str = "train"    # 分词器名称或路径    tokenizer\_name\_or\_path: str = None @dataclass class 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"\\n\\n==============\\n")                    f.write(completion)  # 写入生成的输出            # 定义正则表达式模式,用于匹配 <think> 和 <answer> 标签            regex = r"^<think>(\[^<\]\*(?:<(?!/?think>)\[^<\]\*)\*)<\\/think>\\n<answer>(\[\\s\\S\]\*?)<\\/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 rewards 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):        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"\\n\\n==============\\n")                        f.write(completion)  # 写入生成的输出            else:                rewards.append(0.0)  # 如果不正确,奖励为 0        except Exception:            rewards.append(0.0)  # 如果评估失败,奖励为 0    return rewards def 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 rewards def 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_Datawhale-R1.sh として保存し、ターミナルで bash train_Datawhale-R1.sh を実行してください。

コンピューティングカード番号を制限したい場合は、ここで設定してください(例:cuda:1-3のみを使用する)。制限したくない場合は、以下の行を削除してください。

 export CUDA\_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推論用に1枚のカードを確保する場合、`--num_processes` の値は n-1(使用するコンピューティングカードの数)にする必要があります。例えば、3枚のカードがある場合、`--num_processes` の値は 2 にする必要があります。この値は、`deepspeed_zero3.yaml` で設定されている `num_processes` 値 8 を上書きします。

さらに、前述のように、カスタム ハードウェア構成要件がある場合は、--config_file パラメータを使用しないでください。

このメッセージは、モデルのトレーニングが完了したことを意味します。これで、Swanlab で優れたトレーニングデータを閲覧できます(モバイルでも閲覧可能なので、錬金術師を目指す方に最適です)。

トレーニングプロセスの説明

プロセスの概要

Datawhale-R1 のトレーニング プロセスを確認してみましょう。

  1. プロンプトの単語を Qwen 2.5 モデルに入力します。
  2. Qwen 2.5 は、複数の思慮深い応答を出力します (この実験では 8 に設定されており、num_generations パラメータによって決定されます)。
  3. モデルの回答は 3 つの報酬関数に入力されて計算され、結果が合計されます。
  4. 報酬値は GRPO 戦略に入力され、GRPO は報酬値を使用して Qwen 2.5 モデルを調整する方法を決定します。
  5. 上記のプロセスを繰り返します (この実験は、max_steps パラメータによって決定された 450 回繰り返されました)。

強化学習に馴染みのない方もいらっしゃるかもしれませんので、関連する概念については後続の記事で紹介します。ここでは、例え話をしてみましょう。数学教師(GRPO戦略)とクラス(Qwen 2.5モデル。クラスの生徒は全員同じ能力を持つと仮定)がいる学校を想像してみてください。学校では毎月試験(複数段階)があります。各試験では、生徒は1枚の解答用紙につき1つの質問、1つの質問のみという形で、複数の解答用紙を作成します。(生徒が複数いるため、複数の解答用紙があり、それぞれが複数の思考刺激モデルに対応しています。これらの解答は必ずしも同一ではありません。)その後、数学教師はこれらの解答用紙を採点します(報酬関数の計算)。採点ルールは以下のとおりです。

  1. 回答の形式が正しいかどうかを確認します(フォーマット報酬関数)。
  2. 解は正しいですか(方程式報酬関数)?
  3. 解決手順は十分に詳細ですか (思考長さの報酬関数)?

最後に、各セクションのスコアを合計して、複数のテストスコア(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 class SwanlabArguments:    """SwanLab参数的数据类"""    # 是否使用 SwanLab    swanlab: bool    # SwanLab 用户名    workspace: str    # SwanLab 的项目名    project: str    # SwanLab 的实验名    experiment\_name: str

データホエール-R1.yaml

Swanlabトレーニングプロセスパラメータ記録

swanlab: true # Swanlabを有効にするかどうか
ワークスペース: <ユーザー名>
プロジェクト: <プロジェクト名、複製プロジェクト全体の名前、例: Datawhale-R1-by_xxx>
experiment_name: <実験名、ハイパーパラメータ実行のカスタム名、例: qwen2.5-3B-lr:5e-7_beta:0.001>

次に、grpo_function について説明します。まず、データセットがどのようになっているか確認しましょう。タスクは実は非常にシンプルで、24点ゲームに似ています。例えば、[44, 19, 35] のような数値が与えられた場合、モデルは算術演算を用いて、結果が目標値(例えば 98)と正確に一致する方程式を生成する必要があります。詳細な要件はプロンプトで示します。

プロンプトは以下の通りです。Pythonのf文字列機能を用いて特定の数値を入力し、アシスタントの末尾に`\n<think>`を追加することで、モデルが必要に応じて段階的に思考を開始するように促します。プロンプトは、mini-r1のプロンプトをDeepSeekを用いて翻訳したものです。中国語の読者は中国語を速く読む傾向があります。

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

ここでは、プロンプトをQwen 2.5プロンプトテンプレートに変換し、より使い慣れた方法でプロンプトを受信できるようにします。この変換を段階的に進めていきます。モデルの出力の先頭として `\n<think>` を使用し、書き込みを続けます。サンプルはPython辞書として返します。これにより、TRLは報酬関数を呼び出す際に、対応するパラメータにキーを自動的に割り当てます。さらに、TRLは複数のモデル出力を補完として設定します。

 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 = 98 r1\_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|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im\_end|>\n<|im\_start|>user\n使用给定的数字 [44, 19, 35],创建一个等于98 的方程。你可以使用基本算术运算(+、-、*、/)一次或多次,但每个数字只能使用一次。在 <think> </think> 标签中展示你的思考过程,并在 <answer> </answer> 标签中返回最终方程,例如 <answer> (1 + 2) / 3 </answer>。在 <think> 标签中逐步思考。<|im\_end|>\n<|im\_start|>assistant\n让我们逐步解决这个问题。\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,我们对复现流程的介绍就大致结束了,我们会在文末提供完整的文档,开源我们的复现工作。

训练结果解读

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

对比下面两张图,我们发现模型前期学习输出格式的速度很快,大概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。

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

見通し

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

最后嘱咐一些要点,

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

  • 小于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,并在项目中公布,虽然模型本身没什么用。

复现文件在Datawhale-R1 文件夹。

Unlock-DeepSeekプロジェクトはまだ完了しておらず、急速なイテレーションが進行中です。今後の展開にご期待ください。

いいね (3件のいいね!)↓