DeepSeek-R1 再現プロジェクト: オープンソースで実現する方法

DeepSeek-R1 の完全なオープンソースでの再現を目指す huggingface/open-r1 プロジェクトについて解説します。この記事では、プロジェクトの概要、インストール方法、モデルのトレーニング、評価、データ生成の手順、そしてコントリビューションの方法を詳しく解説します。

Open R1 プロジェクトとは？

このプロジェクトの目的は、DeepSeek-R1 のパイプラインを再現し、誰もがその上に構築できるように、不足している要素を構築することです。設計はシンプルで、主に以下の要素で構成されています。

• src/open_r1: モデルのトレーニングと評価、および合成データを生成するためのスクリプトが含まれています。
• Makefile: 上記のスクリプトを利用して、R1 パイプラインの各ステップを簡単に実行できるコマンドが含まれています。

DeepSeek-R1 の再現ステップ

DeepSeek-R1 のテクニカルレポートを参考に、以下の3つの主要なステップで進めます。

1. R1-Distill モデルの再現: DeepSeek-R1 から高品質なコーパスを蒸留します。
2. 純粋な RL パイプラインの再現: DeepSeek が R1-Zero を作成するために使用したパイプラインを再現します。数学、推論、コーディングのための大規模なデータセットをキュレーションする必要があります。
3. ベースモデルから RL チューニングへの移行: 複数段階のトレーニングを通じて、ベースモデルから RL チューニングされたモデルへの移行を可能にすることを示します。

インストール手順

コードを実行するには、まず Conda などのツールを使用して Python 仮想環境を作成します。

conda create -n openr1 python=3.11 && conda activate openr1

次に、vLLM をインストールします。

pip install vllm==0.6.6.post1

# HF の場合 (クラスターには CUDA 12.1 のみがあります)
pip install vllm==0.6.6.post1 --extra-index-url https://download.pytorch.org/whl/cu121

PyTorch v2.5.1 がインストールされることに注意してください。vLLM バイナリはこのバージョン用にコンパイルされているため、このバージョンを使用することは非常に重要です。その後、pip install -e .[LIST OF MODES] を使用して、特定のユースケースに必要な残りの依存関係をインストールできます。

貢献者のほとんどには、以下をお勧めします。

pip install -e ".[dev]"

次に、Hugging Face および Weights and Biases のアカウントにログインします。

huggingface-cli login
wandb login

最後に、モデル/データセットを Hugging Face Hub にロードおよびプッシュできるように、システムに Git LFS がインストールされていることを確認します。

git-lfs --version

インストールされていない場合は、以下を実行します。

sudo apt-get install git-lfs

モデルのトレーニング

DDP または DeepSpeed ZeRO-2 および ZeRO-3 のいずれかを使用してモデルをトレーニングできます。メソッドを切り替えるには、configs の accelerate YAML 設定へのパスを変更します。

SFT（Supervised Fine-Tuning）

DeepSeek-R1 から蒸留された、推論のトレースを持つデータセット (例: Bespoke-Stratos-17k) で SFT を実行するには、以下を実行します。

accelerate launch --config_file=configs/zero3.yaml src/open_r1/sft.py \
 --model_name_or_path Qwen/Qwen2.5-Math-1.5B-Instruct \
 --dataset_name HuggingFaceH4/Bespoke-Stratos-17k \
 --learning_rate 2.0e-5 \
 --num_train_epochs 1 \
 --packing \
 --max_seq_length 4096 \
 --per_device_train_batch_size 4 \
 --per_device_eval_batch_size 4 \
 --gradient_accumulation_steps 4 \
 --gradient_checkpointing \
 --bf16 \
 --logging_steps 5 \
 --eval_strategy steps \
 --eval_steps 100 \
 --output_dir data/Qwen2.5-1.5B-Open-R1-Distill

Slurm ジョブを起動するには、以下を実行します。

sbatch --output=/path/to/logs/%x-%j.out --err=/path/to/logs/%x-%j.err slurm/sft.slurm {model} {dataset} {accelerator}

ここで {model} および {dataset} は、Hugging Face Hub のモデルおよびデータセット ID を指し、{accelerator} は configs の 🤗 Accelerate 設定の選択を指します。

GRPO（Generative Reweighted Policy Optimization）

accelerate launch --config_file configs/zero3.yaml src/open_r1/grpo.py \
 --output_dir DeepSeek-R1-Distill-Qwen-7B-GRPO \
 --model_name_or_path deepseek-ai/DeepSeek-R1-Distill-Qwen-7B \
 --dataset_name AI-MO/NuminaMath-TIR \
 --max_prompt_length 256 \
 --per_device_train_batch_size 1 \
 --gradient_accumulation_steps 16 \
 --logging_steps 10 \
 --bf16

モデルの評価

lighteval を使用してモデルを評価します。カスタムタスクは src/open_r1/evaluate.py で定義されています。単一の GPU に収まるモデルの場合は、以下を実行します。

MODEL=deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
MODEL_ARGS="pretrained=$MODEL,dtype=float16,max_model_length=32768,gpu_memory_utilisation=0.8"
TASK=aime24
OUTPUT_DIR=data/evals/$MODEL

lighteval vllm $MODEL_ARGS "custom|$TASK|0|0" \
 --custom-tasks src/open_r1/evaluate.py \
 --use-chat-template \
 --system-prompt="Please reason step by step, and put your final answer within \boxed{}." \
 --output-dir $OUTPUT_DIR

複数の GPU でスループットを向上させるには、以下のように データ並列処理 を使用します。

NUM_GPUS=8
MODEL=deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B
MODEL_ARGS="pretrained=$MODEL,dtype=float16,data_parallel_size=$NUM_GPUS,max_model_length=32768,gpu_memory_utilisation=0.8"
TASK=aime24
OUTPUT_DIR=data/evals/$MODEL

lighteval vllm $MODEL_ARGS "custom|$TASK|0|0" \
 --custom-tasks src/open_r1/evaluate.py \
 --use-chat-template \
 --system-prompt="Please reason step by step, and put your final answer within \boxed{}." \
 --output-dir $OUTPUT_DIR

大きなモデルを GPU 間でシャーディングする必要がある場合は、テンソル並列処理 を使用して以下を実行します。

NUM_GPUS=8
MODEL=deepseek-ai/DeepSeek-R1-Distill-Qwen-32B
MODEL_ARGS="pretrained=$MODEL,dtype=float16,tensor_parallel_size=$NUM_GPUS,max_model_length=32768,gpu_memory_utilisation=0.8"
TASK=aime24
OUTPUT_DIR=data/evals/$MODEL

export VLLM_WORKER_MULTIPROC_METHOD=spawn
lighteval vllm $MODEL_ARGS "custom|$TASK|0|0" \
 --custom-tasks src/open_r1/evaluate.py \
 --use-chat-template \
 --system-prompt="Please reason step by step, and put your final answer within \boxed{}." \
 --output-dir $OUTPUT_DIR

モデル、タスク、およびオプションで並列処理技術と GPU の数を指定して、make evaluate で評価を開始することもできます。

データ生成

Distilabel を使用したデータ生成

distilabel を使用して、DeepSeek-R1 モデルからデータを生成することができます。

小さな蒸留された R1 モデルからデータを生成する場合

次の例は、1xH100 で実行できます。まず、次の依存関係をインストールします。

pip install "distilabel[vllm]>=1.5.2"

次に、次のスニペットを pipeline.py という名前のファイルに保存し、python pipeline.py で実行します。これにより、10 個の例それぞれに対して 4 つの生成が生成されます (リポジトリのユーザー名を組織/ユーザー名に変更します)。

from datasets import load_dataset
from distilabel.models import vLLM
from distilabel.pipeline import Pipeline
from distilabel.steps.tasks import TextGeneration

prompt_template = """\
You will be given a problem. Please reason step by step, and put your final answer within \boxed{}:
{{ instruction }}"""

dataset = load_dataset("AI-MO/NuminaMath-TIR", split="train").select(range(10))

model_id = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B" # Exchange with another smol distilled r1

with Pipeline(
 name="distill-qwen-7b-r1",
 description="A pipeline to generate data from a distilled r1 model",
) as pipeline:

 llm = vLLM(
 model=model_id,
 tokenizer=model_id,
 extra_kwargs={
 "tensor_parallel_size": 1,
 "max_model_len": 8192,
 },
 generation_kwargs={
 "temperature": 0.6,
 "max_new_tokens": 8192,
 },
 )
 prompt_column = "problem"
 text_generation = TextGeneration(
 llm=llm,
 template=prompt_template,
 num_generations=4,
 input_mappings={"instruction": prompt_column} if prompt_column is not None else {}
 )

if __name__ == "__main__":
 distiset = pipeline.run(dataset=dataset)
 distiset.push_to_hub(repo_id="username/numina-deepseek-r1-qwen-7b")

DeepSeek-R1 からデータを生成する場合

より大きな DeepSeek-R1 を実行するために、各 8×H100 GPU を備えた 2 つのノードを、このリポジトリの slurm/generate.slurm にある slurm ファイルを使用して使用しました。まず、依存関係をインストールします。

pip install https://wheels.vllm.ai/221d388cc5a836fa189305785ed7e887cea8b510/vllm-1.0.0.dev-cp38-abi3-manylinux1_x86_64.whl --extra-index-url https://download.pytorch.org/whl/cu121

pip install "distilabel[vllm,ray,openai]>=1.5.2"

次に、次のコマンドを実行します。

sbatch slurm/generate.slurm \
 --hf-dataset AI-MO/NuminaMath-TIR \
 --temperature 0.6 \
 --prompt-column problem \
 --model deepseek-ai/DeepSeek-R1 \
 --hf-output-dataset username/r1-dataset

コントリビューション

コントリビューションは大歓迎です。詳しくは #23 を参照してください。

まとめ

この記事では、DeepSeek-R1 のオープンソースでの再現を目指す huggingface/open-r1 プロジェクトについて詳しく解説しました。このプロジェクトに参加し、DeepSeek-R1 の再現と発展に貢献しましょう。