本文使用PyTorch从零构建 DeepSeek R1 模型,详解模型架构与训练流程,深入了解强化学习、拒绝采样等关键技术。
原文标题:用PyTorch从零构建 DeepSeek R1:模型架构和分步训练详解
原文作者:数据派THU
冷月清谈:
怜星夜思:
2、文章中提到了多种奖励函数,你认为哪个奖励函数对 DeepSeek R1 模型最终的推理能力影响最大?为什么?如果可以,你会如何改进这些奖励函数?
3、在冷启动数据准备阶段,文章提到了少样本 Prompting、直接 Prompting 和后处理优化等方法,你认为哪种方法对提升数据质量最有效?如果你要构建自己的推理数据集,你会如何综合运用这些方法?
原文内容
来源:Deephub IMBA本文约1.8万字,建议阅读15+分钟本文深入剖析了 DeepSeek R1 模型的构建过程。
train-deepseek-r1/ ├── code.ipynb # Jupyter Notebook 代码实现 ├── requirements.txt # 依赖库列表 └── r1_for_dummies.md # 面向非技术受众的 DeepSeek R1 解释
环境配置
git clone https://github.com/FareedKhan-dev/train-deepseek-r1.git cd train-deepseek-r1 pip install -r requirements.txt
# 导入必要的库 import logging import os import sys import re import math from dataclasses import dataclass, field from typing import List, Optional
导入 PyTorch 与deep hub Hugging Face Transformers
import torch
import transformers
from transformers import (
AutoModelForCausalLM,
AutoTokenizer,
HfArgumentParser,
TrainingArguments,
set_seed,
TrainerCallback,
TrainerControl,
TrainerState,
)
from transformers.trainer_utils import get_last_checkpoint导入数据集工具库
import datasets
from datasets import load_dataset导入 TRL (Transformers Reinforcement Learning deep—hub) 库
from trl import (
AutoModelForCausalLMWithValueHead,
PPOConfig,
PPOTrainer,
GRPOTrainer,
GRPOConfig,
SFTTrainer
)导入数学相关工具库
from latex2sympy2_extended import NormalizationConfig
from math_verify import LatexExtractionConfig, parse, verify
训练数据集
-
NuminaMath-TIR:用于 R1 Zero 阶段的训练。
-
Bespoke-Stratos-17k:用于 R1 阶段的训练。
#### 输出 ####
{
'problem': 'What is the degree of the polynomial 4 +5x^3 ... ',
'solution': 'This polynomial is not written in ...',
'messages': [{'from': 'user', 'value': 'The problem ...'}]
}
#### 输出 ####
#### 输出 ####
{
'system': 'Your role as an assistant involves ... ',
'conversations': [{'from': 'user', 'value': 'Return your ... deep hub'}]
}
##### 输出 ####
DeepSeek R1 训练流程概览
以上是对 DeepSeek R1 训练流程的高度概括,后续章节将深入剖析每个训练阶段的具体细节。
基础模型选型
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" OUTPUT_DIR = "data/Qwen-GRPO-training" # 用于保存训练后模型
创建输出目录,如果目录不存在
os.makedirs(OUTPUT_DIR, exist_ok=True)
初始化 tokenizer,并指定聊天模板
tokenizer = AutoTokenizer.from_pretrained(
MODEL_NAME,
trust_remote_code=True,
padding_side=“right”
)若 pad token 未设置deephub,则指定 pad token 为 eos token
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
print(f"Vocabulary size: {len(tokenizer)}“)
print(f"Model max length: {tokenizer.model_max_length}”)
print(f"Pad token: {tokenizer.pad_token}“)
print(f"EOS token: {tokenizer.eos_token}”)
#### 输出 ####
Vocabulary size: 151665
Model max length: 131072
Pad token: <|endoftext|>
EOS token: <|im_end|>
#### 输出 ####
# 初始化基础模型 model = AutoModelForCausalLM.from_pretrained( MODEL_NAME, trust_remote_code=True, torch_dtype=torch.bfloat16 )
print(f"Model parameters: {model.num_parameters():,}")
#### 输出 ####
Model parameters: 494,032,768
#### 输出 ####
# 检查 CUDA 是否可用 device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"Using device: {device}")
将模型移至可用设备
model.to(device)
测试基础推理能力
def test_model_inference(user_input: str):
“”“使用已加载的模型和 tokenizer 测试基础模型推理。”“”
messages = [
{“role”: “system”, “content”: “You are Qwen, a helpful assistant.”},
{“role”: “user”, “content”: user_input}
]应用聊天模板
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)分词并生成
inputs = tokenizer(text, return_tensors=“pt”).to(device)
outputs = model.generate(
**inputs,
max_new_tokens=100,
do_sample=True,
temperature=0.7
)response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return response测试模型
test_input = “how are you?”
response = test_model_inference(test_input)
print(f"Test Input: {test_input}“)
print(f"Model Response: {response}”)
#### 输出 ####
"Test Input: how are you?
Model Response: As an AI language model I dont have feelings ..."
##### 输出 ####
强化学习 (RL) 框架中的策略模型 (R)
R1 Zero 的 GRPO 算法
-
准确性 (Accuracy):答案是否在数学上正确?
-
格式 (Format):是否规范使用了 <think> 和 <answer> 标签?
-
推理步骤 (Reasoning Steps):推理逻辑是否清晰可循?
-
余弦缩放 (Cosine Scaling):响应内容是否精炼简洁?
-
重复惩罚 (Repetition Penalty):是否存在不必要的重复内容?
Prompt 模板
# 基于 GRPO 训练的 DeepSeek 系统 Prompt SYSTEM_PROMPT = ( "A conversation between User and Assistant. The user asks a question, \ and the Assistant solves it. The assistant " "first thinks about the reasoning process in the mind and \ then deephub provides the user with the answer. The reasoning " "process and answer are enclosed within <think> </think> \ and <answer> </answer> tags, respectively, i.e., " "<think> reasoning process here </think><answer> answer here </answer>" )
训练数据预处理
# 构建训练数据结构的函数 def make_conversation(example): """将数据集样本转换为对话格式。""" return { "prompt": [ {"role": "system", "content": SYSTEM_PROMPT}, {"role": "user", "content": example["problem"]}, ], }
# 加载并准备数据集 def load_math_dataset(): """加载并准备数学数据集。""" dataset = load_dataset( "AI-MO/NuminaMath-TIR", name="default", split=['train', 'test'] )
将数据集划分为训练集和测试集
dataset = {
‘train’: dataset[0],
‘test’: dataset[1]
}应用对话格式
for split in dataset:
dataset[split] = dataset[split].map(make_conversation)若存在 ‘messages’ 列,则移除该列
if “messages” in dataset[split].column_names:
dataset[split] = dataset[split].remove_columns(“messages”)
return dataset
# 加载训练数据集并打印训练集/测试集大小 dataset = load_math_dataset()
print(f"Train set size: {len(dataset[‘train’])}“)
print(f"Test set size: {len(dataset[‘test’])}”)
#### 输出 ####
Train set size: 72441
Test set size: 99
#### 输出 ####
def validate_dataset(dataset): """对数据集执行基础验证检查。"""
定义数据集所需字段
required_fields = [“problem”, “prompt”]
遍历数据集的 ‘train’ 和 ‘test’ 划分
for split in [‘train’, ‘test’]:
print(f"\nValidating {split} split:")从数据集中获取列名
fields = dataset[split].column_names
检查是否缺少必要字段
missing = [field for field in required_fields if field not in fields]
if missing:
print(f"Warning: Missing fields: {missing}") # 若缺少字段,则发出警告
else:
print(“✓ All required fields present”) # 确认所有必要字段均存在获取数据集划分的首个样本
sample = dataset[split][0]
提取包含对话消息列表的 ‘prompt’ 字段
messages = sample[‘prompt’]
验证 Prompt 格式:
- 至少包含两条消息
- 首条消息 Role 为 ‘system’
- 次条消息 Role 为 ‘user’
if (len(messages) >= 2 and
messages[0][‘role’] == ‘system’ and
messages[1][‘role’] == ‘user’):
print(“✓ Prompt format is correct”) # 确认 Prompt 格式正确
else:
print(“Warning: Incorrect prompt format”) # 若 Prompt 格式不正确,则发出警告验证数据集
validate_dataset(dataset)
Validating train split:
✓ All required fields present
✓ Prompt format is correctValidating test split:
✓ All required fields present
✓ Prompt format is correct
奖励函数
-
准确性 (Accuracy):答案是否在数学上正确?
-
格式 (Format):是否规范使用了 <think> 和 <answer> 标签?
-
推理步骤 (Reasoning Steps):推理逻辑是否清晰可循?
-
余弦缩放 (Cosine Scaling):响应内容是否精炼简洁?
-
重复惩罚 (Repetition Penalty):是否存在不必要的重复内容?
准确性奖励函数
def accuracy_reward(completions, solution, **kwargs): """ 奖励函数,用于检查模型的响应是否在数学上等价于标准答案。 使用 deep hub latex2sympy2 进行解析,使用 math_verify 进行验证。 """
提取模型响应内容
contents = [completion[0][“content”] for completion in completions]
rewards =for content, sol in zip(contents, solution):
解析标准答案
gold_parsed = parse(sol, extraction_mode=“first_match”,
extraction_config=[LatexExtractionConfig()])if gold_parsed: # 检查标准答案是否解析成功
使用宽松的归一化配置解析模型答案
answer_parsed = parse(
content,
extraction_config=[
LatexExtractionConfig(
normalization_config=NormalizationConfig(
nits=False,
malformed_operators=False,
basic_latex=True,
equations=True,
boxed=“all”,
units=True,
),
boxed_match_priority=0,
try_extract_without_anchor=False,
)
],
extraction_mode=“first_match”,
)若答案正确,奖励 1.0,否则奖励 0.0
reward = float(verify(answer_parsed, gold_parsed))
else:若标准答案解析失败,则给予中性奖励 0.5
reward = 0.5
print(“Warning: Failed to parse gold solution:”, sol) # 警告:无法解析标准答案rewards.append(reward)
return rewards
-
使用 latex2sympy2 工具将标准答案转换为结构化的数学表达式。
-
若标准答案解析失败,则给予中性奖励 0.5。
-
提取模型输出,并进行归一化处理,以提高评估的鲁棒性。
-
利用 math_verify 工具,比对解析后的模型响应与解析后的标准答案是否在数学上一致。
-
若数学上一致,则奖励 1,否则奖励 0。
格式奖励函数
# 实现格式奖励函数 def format_reward(completions, **kwargs): """ 奖励函数,用于检查模型输出是否符合预定义的格式: <think>...</think>deep hub <answer>...</answer>。 """
定义目标格式的正则表达式模式
pattern = r"^<think>.?</think>\s<answer>.*?</answer>$"
从每个模型输出中提取内容
completion_contents = [completion[0][“content”] for completion in completions]
检查每个模型输出是否与目标模式匹配
matches = [re.match(pattern, content, re.DOTALL | re.MULTILINE)
for content in completion_contents]若格式正确,奖励 1.0,否则奖励 0.0
return [1.0 if match else 0.0 for match in matches]
-
定义正则表达式 (regex) 模式。该模式精确描述了期望的输出格式:以 <think> 开头,<think> 和 </think> 标签对之间可包含任意字符,随后是空白字符,然后以 <answer> 开头,<answer> 和 </answer> 标签对之间可包含任意字符,并以此结尾。
-
从每个模型的输出结果中提取文本内容。
-
使用 re.match 函数,逐一检查模型输出内容是否与定义的正则表达式模式完全匹配。re.DOTALL 标志使正则表达式中的 . 能够匹配换行符,re.MULTILINE 标志使 ^ 和 $ 能够匹配整个字符串的起始和结束位置,而非仅限于行首和行尾。
-
对于符合格式规范的模型输出,奖励值设为 1;反之,设为 0。该奖励机制对格式的正确性采取严格的二元评价标准。
推理步骤奖励函数
-
“步骤 1”、“步骤 2” 等序号型步骤标识;
-
“1.”、“2.” 等数字编号列表;
-
“-” 或 “*” 等项目符号列表;
-
“首先”、“其次”、“然后”、“最后” 等过渡性连接词。
def reasoning_steps_reward(completions, **kwargs): r""" 奖励函数,用于鼓励模型进行清晰的逐步推理。 该函数会检测诸如 "Step 1:"、编号列表、项目符号以及过渡词等模式。 """
用于匹配推理步骤指示符的正则表达式模式
pattern = r"(Step \d+:|^\d+.|\n-|\n*|First,|Second,|Next,|Finally,)"
提取模型输出内容
completion_contents = [completion[0][“content”] for completion in completions]
统计每个模型输出中推理步骤指示符的数量
matches = [len(re.findall(pattern, content, re.MULTILINE))
for content in completion_contents]奖励值与推理步骤数量成正比,最高奖励值为 1.0
此处采用“魔法数字” 3,鼓励模型至少输出 3 个推理步骤以获得全额奖励
return [min(1.0, count / 3) for count in matches]
-
若模型输出中包含 3 个或更多推理指示符(count >= 3),则奖励值为 1.0(满额奖励)。
-
若指示符数量少于 3 个(例如,count = 1 或 2),则获得 部分 奖励(例如,1/3 或 2/3)。
-
若未检测到任何推理指示符(count = 0),则奖励值为 0.0。
余弦缩放奖励函数
-
对于正确答案:我们更倾向于奖励 简洁、直接的解答,而非冗长、散漫的答案。简明扼要且正确的答案通常更佳。
-
对于错误答案: 相较于尝试进行推理的较长错误答案,简短的错误答案可能更不可取。因此,我们希望对简短的错误答案施加 更重 的惩罚,而对较长的错误答案施加相对较轻的惩罚。
# 实现余弦缩放奖励函数 def get_cosine_scaled_reward( min_value_wrong: float = -0.5, max_value_wrong: float = -0.1, min_value_correct: float = 0.8, max_value_correct: float = 1.0, max_len: int = 1000, ): """ 返回一个余弦缩放奖励函数。该函数基于模型输出的长度,对准确性奖励进行缩放调整。 较短的正确答案将获得更高的奖励,而较长的错误答案将受到较轻的惩罚。 """ def cosine_scaled_reward(completions, solution, accuracy_rewards, **kwargs): """ 余弦缩放奖励函数,根据模型输出长度调整准确性奖励。 """ contents = [completion[0]["content"] for completion in completions] rewards = []
for content, sol, acc_reward in zip(contents, solution, accuracy_rewards):
gen_len = len(content) # 模型生成答案的长度
progress = gen_len / max_len # 答案长度相对于最大长度的进度
cosine = math.cos(progress * math.pi) # 基于进度的余弦值if acc_reward > 0.5: # 假设准确性奖励函数对正确答案给出约 1.0 的奖励
min_value = min_value_correct
max_value = max_value_correct
else: # 答案错误
min_value = max_value_wrong # 注意此处交换了 min_value 和 max_value
max_value = min_value_wrong余弦缩放公式
reward = min_value + 0.5 * (max_value - min_value) * (1.0 + cosine)
rewards.append(float(reward))
return rewards
return cosine_scaled_reward
重复惩罚奖励函数
def get_repetition_penalty_reward(ngram_size: int = 3, max_penalty: float = -0.1): """ 返回一个重复惩罚奖励函数。该函数惩罚模型在生成文本中对 n-gram 的重复使用。 """ if max_penalty > 0: raise ValueError(f"max_penalty {max_penalty} should not be positive")
def zipngram(text: str, ngram_size: int):
“”“辅助函数,用于从文本中生成 n-gram。”“”
words = text.lower().split() # 转换为小写并按空格分割为单词列表
return zip(*[words[i:] for i in range(ngram_size)]) # 生成 n-gramdef repetition_penalty_reward(completions, **kwargs) -> float:
“”"
重复惩罚奖励函数。
“”"
contents = [completion[0][“content”] for completion in completions]
rewards =
for completion in contents:
if completion == “”: # 对于空输出,不施加惩罚
rewards.append(0.0)
continue
if len(completion.split()) < ngram_size: # 对于过短的输出,不施加惩罚
rewards.append(0.0)
continuengrams = set() # 使用集合存储唯一的 n-gram
total = 0
for ng in zipngram(completion, ngram_size): # 生成 n-gram
ngrams.add(ng) # 将 n-gram 添加到集合 (重复的 n-gram 会被自动忽略)
total += 1 # 统计 n-gram 的总数量计算缩放因子:重复程度越高 -> 缩放因子越大
scaling = 1 - len(ngrams) / total
reward = scaling * max_penalty # 基于缩放因子施加惩罚
rewards.append(reward)
return rewards
return get_repetition_penalty_reward
至此,五个奖励函数已全部实现。接下来,我们将进入配置训练参数的环节。
R1 Zero 训练配置
# 为 GRPO 脚本参数定义 GRPOScriptArguments 类,用于配置奖励函数参数 @dataclass class GRPOScriptArguments: """ GRPO 训练的脚本参数,特别是与奖励函数相关的参数。 """
reward_funcs: list[str] = field(
default_factory=lambda: [“accuracy”, “format”],
metadata={
“help”: “奖励函数列表。可选值: ‘accuracy’, ‘format’, ‘reasoning_steps’, ‘cosine’, ‘repetition_penalty’”
},
)
cosine_min_value_wrong: float = field(
default=-0.5,
metadata={“help”: “余弦缩放奖励函数中,错误答案的最小奖励值”},
)
cosine_max_value_wrong: float = field(
default=-0.1,
metadata={“help”: “余弦缩放奖励函数中,错误答案的最大奖励值”},
)
cosine_min_value_correct: float = field(
default=0.8,
metadata={“help”: “余弦缩放奖励函数中,正确答案的最小奖励值”},
)
cosine_max_value_correct: float = field(
default=1.0,
metadata={“help”: “余弦缩放奖励函数中,正确答案的最大奖励值”},
)
cosine_max_len: int = field(
default=1000,
metadata={“help”: “余弦缩放奖励函数的最大长度阈值”},
)
repetition_n_grams: int = field(
default=3,
metadata={“help”: “重复惩罚奖励函数中,n-gram 的大小”},
)
repetition_max_penalty: float = field(
default=-0.1,
metadata={“help”: “重复惩罚奖励函数中,最大惩罚值 (负值)”},
)
# 从 transformers 库定义 TrainingArguments training_args = TrainingArguments( output_dir=OUTPUT_DIR, # 检查点和日志输出目录 overwrite_output_dir=True, num_train_epochs=1, # 训练的总 epoch 数 per_device_train_batch_size=8, # 每个设备的训练批次大小 per_device_eval_batch_size=16, # 评估批次大小 gradient_accumulation_steps=2, # 梯度累积步数,用于模拟更大的批次大小 learning_rate=5e-5, # AdamW 优化器的初始学习率 warmup_ratio=0.1, # 预热步数比例 weight_decay=0.01, # 权重衰减系数,应用于除 bias 和 LayerNorm 权重外的所有层 logging_steps=10, # 日志记录频率 (步数) evaluation_strategy="steps", # 评估策略:每 `eval_steps` 步进行评估 eval_steps=50, # 评估频率 (步数) save_strategy="steps", # 模型保存策略:每 `save_steps` 步保存模型 save_steps=50, # 模型保存频率 (步数) save_total_limit=2, # 最大 checkpoint 保存数量,超出限制则删除旧 checkpoint dataloader_num_workers=2, # 数据加载器 worker 数量 seed=42, # 随机种子,用于保证实验可复现 bf16=True, # 启用混合精度 BF16 训练 push_to_hub=False, # 是否将模型推送至 Hugging Face Hub gradient_checkpointing=True, # 启用梯度检查点 report_to="none", # 不使用任何报告工具 )
@dataclass class ModelConfig: """ 模型配置类。 """ model_name_or_path: str = field( default=MODEL_NAME, metadata={"help": "预训练模型路径或 Hugging Face Model Hub 模型标识符"} ) model_revision: Optional[str] = field( default="main", metadata={"help": "指定模型版本 (分支名, tag 名 或 commit id)"} ) torch_dtype: Optional[str] = field( default="bfloat16", metadata={"help": "覆盖默认 torch_dtype,以指定 dtype 加载模型"} ) trust_remote_code: bool = field( default=True, metadata={"help": "加载模型和 tokenizer 时,信任远程代码"} ) attn_implementation: Optional[str] = field( default="flash_attention_2", metadata={"help": "选择 Attention 实现方式, 可选 'flash_attention_2' 或 None"} )
# 实例化配置对象 script_args = GRPOScriptArguments() model_args = ModelConfig()
# 实用函数,根据脚本参数获取奖励函数列表 def get_reward_functions(script_args): """ 根据脚本参数,返回奖励函数列表。 """ reward_funcs_list = [] reward_funcs_registry = { "accuracy": accuracy_reward, # 假设 accuracy_reward 函数已在之前步骤定义 "format": format_reward, # 假设 format_reward 函数已在之前步骤定义 "reasoning_steps": reasoning_steps_reward, # 假设 reasoning_steps_reward 函数已定义 "cosine": get_cosine_scaled_reward( # 假设 get_cosine_scaled_reward 函数已定义 min_value_wrong=script_args.cosine_min_value_wrong, max_value_wrong=script_args.cosine_max_value_wrong, min_value_correct=script_args.cosine_min_value_correct, max_value_correct=script_args.cosine_max_value_correct, max_len=script_args.cosine_max_len, ), "repetition_penalty": get_repetition_penalty_reward( # 假设 get_repetition_penalty_reward 函数已定义 ngram_size=script_args.repetition_n_grams, max_penalty=script_args.repetition_max_penalty, ), }
for func_name in script_args.reward_funcs:
if func_name not in reward_funcs_registry:
raise ValueError(f"Reward function ‘{func_name}’ not found in registry.")
reward_funcs_list.append(reward_funcs_registry[func_name])
return reward_funcs_list
logger = logging.getLogger(__name__) class LoggingCallback(TrainerCallback): """ 一个简单的回调函数,用于在特定步骤记录训练信息。 """ def on_step_end(self, args: TrainingArguments, state: TrainerState, control: TrainerControl, **kwargs): if state.global_step % args.logging_steps == 0: logger.info(f"Step {state.global_step}: Loss = {state.log_history[-1].get('loss', None)}, Learning Rate = {state.log_history[-1].get('learning_rate', None)}")
def get_callbacks(training_args, model_args, script_args):
“”"
返回训练过程中使用的回调函数列表。
目前仅包含 LoggingCallback。您可以扩展此列表以添加更多回调函数。
“”"
callbacks = [LoggingCallback()] # 实例化 LoggingCallback
return callbacks
# 获取奖励函数和回调函数 reward_functions = get_reward_functions(script_args) callbacks = get_callbacks(training_args, model_args, script_args)
GRPO 训练循环
# 从 TrainingArguments 创建 GRPOConfig grpo_config = GRPOConfig( **training_args.to_dict(), # 将 TrainingArguments 转换为字典并解包 **{ # 此处移除了 model_init_kwargs # 因为我们直接传递了实例化的 'model' 对象,GRPOTrainer 无需 model_init_kwargs } )
grpo_trainer = GRPOTrainer(
model=model, # 初始化的 Qwen 模型
reward_funcs=reward_functions, # 前述步骤定义的奖励函数列表
args=grpo_config, # GRPOConfig 对象 (由 TrainingArguments 创建)
train_dataset=dataset[‘train’], # 训练数据集
eval_dataset=dataset[‘test’], # 评估数据集
callbacks=callbacks # 回调函数列表
)
# 启动 GRPO 训练循环 train_result = grpo_trainer.train()
...
INFO:dee phub__main__:Step 10: Loss = …, Learning Rate = …
INFO:deeph ub main:Step 20: Loss = …, Learning Rate = …
…
保存 Tiny R1 Zero LLM
# 加载 tokenizer - 如有需要,请确保设置 trust_remote_code=True tokenizer = AutoTokenizer.from_pretrained( TRAINED_MODEL_PATH, trust_remote_code=True, # 如果模型配置需要,则设置为 True padding_side="right" # 确保 padding 方向一致 )
若 pad token 未正确保存或加载,则进行设置
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token加载训练后的模型
trained_model = AutoModelForCausalLM.from_pretrained(
TRAINED_MODEL_PATH,
trust_remote_code=True, # 如果模型架构需要,则设置为 True
torch_dtype=torch.bfloat16 # 保持与训练时一致的数据类型
)将加载的模型移至指定设备 (如有 GPU 可用,则移至 GPU)
trained_model.to(device) # ‘device’ 变量仍为之前定义的 CUDA 设备
# 使用训练后的模型进行推理测试 def test_trained_model_inference(user_input: str): """使用加载的训练后模型和 tokenizer 进行推理测试。""" messages = [ {"role": "system", "content": SYSTEM_PROMPT}, # 复用之前的系统 Prompt {"role": "user", "content": user_input} ]
使用 tokenizer 应用聊天模板
text = tokenizer.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=True
)对输入文本进行分词
inputs = tokenizer(text, return_tensors=“pt”).to(device)
使用 trained_model 生成模型输出
outputs = trained_model.generate(
**inputs,
max_new_tokens=200, # 相比之前,可以生成稍长的文本
do_sample=True,
temperature=0.7
)将生成的 token 解码为文本
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return response
R1 Zero 的主要问题
-
“你 必须 使用逐步推理”(仅使用了 “推理过程” 这一宽泛表述,将推理方式的定义权交由模型自身)。
-
“你 必须 使用反思性推理”。
-
“你 必须 采用特定的问题解决策略”。
正是上述两个主要问题,促使 DeepSeek 团队将初始的 R1 Zero 模型迭代升级为 R1 模型。
为 SFT 准备冷启动数据
基于长 CoT 的少样本 Prompting
# 加载模型和 Tokenizer MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True, padding_side="right") if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, trust_remote_code=True, torch_dtype=torch.bfloat16).to("cuda" if torch.cuda.is_available() else "cpu")
生成长 CoT 响应
def generate_response(prompt_text):
messages = [
{“role”: “system”, “content”: “dee phub You are a helpful assistant that provides step-by-step solutions.”},
{“role”: “user”, “content”: prompt_text}
]
text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
inputs = tokenizer(text, return_tensors=“pt”).to(model.device)
outputs = model.generate(**inputs, max_new_tokens=200, do_sample=False) # 为示例保持确定性输出
response = tokenizer.decode(outputs[0], skip_special_tokens=True)
return response.split(“<|im_start|>assistant\n”)[-1].strip() # 提取助手的响应
# 示例问题及其解答 (使用 | special_token | 作为分隔符) few_shot_prompt = """ Problem: dee ph ub What's the square root of 9 plus 5? Solution: <|special_token|> First, find the square root of 9, which is 3. Then, add 5 to 3. 3 + 5 equals 8. <|special_token|> Summary: The answer is 8.
Problem: Train travels at 60 mph for 2 hours, how far?
Solution: <|special_token|> Use the formula: Distance = Speed times Time. Speed is 60 mph, Time is 2 hours. Distance = 60 * 2 = 120 miles. <|special_token|> Summary: Train travels 120 miles.Problem: What is 2 + 3 * 4?
Solution:
“”"
使用基础模型进行少样本生成:
# 使用少样本示例,为目标问题生成响应 target_problem_prompt = few_shot_prompt + "What is 2 + 3 * 4?" model_response_few_shot = generate_response(target_problem_prompt)
print(“Few-shot Prompt:”)
print(target_problem_prompt)
print(“\nModel Response (Few-shot CoT):”)
print(model_response_few_shot)
Few-shot Prompt: Problem: What's the square root of 9 plus 5? Solution: <|special_token|> First, find the square root of 9, which is 3. Then, add 5 to 3. 3 + 5 equals 8. <|special_token|> Summary: The answer is 8.
Problem: Train travels at 60 mph for 2 hours, how far?
Solution: <|special_token|> Use the formula: Distance = Speed times Time.
Speed is 60 mph, Time is 2 hours. Distance = 60 * 2 = 120 miles.
<|special_token|> Summary: Train travels 120 miles.Problem: What is 2 + 3 * 4?
Solution:
Model Response (Few-shot CoT):
<|special_token|> To solve 2 + 3 * 4, we need to follow the order
of operations (PEMDAS/BODMAS). Multiplication should be performed
before addition.
Step 1: Multiply 3 by 4, which equals 12.
Step 2: Add 2 to the result from Step 1: 2 + 12 = 14.
<|special_token|> Summary: The answer is 14.
直接 Prompting
# 直接 prompting 示例 direct_prompt_text = """ Problem: d ee p hub Solve this, show reasoning step-by-step, and verify: What is 2 + 3 * 4? """
model_response_direct = generate_response(direct_prompt_text)
print(“Direct Prompt:”)
print(direct_prompt_text)
print(“\nModel Response (Direct Prompting):”)
print(model_response_direct)
Direct Prompt: Problem: Solve this, show reasoning step-by-step, and verify: What is 2 + 3 * 4?
Model Response (Direct Prompting):
<|special_token|> Reasoning: To solve 2 + 3 * 4, I need to follow
the order of operations, which states that multiplication should
be done before addition.
Step 1: Multiply 3 by 4, which equals 12.
Step 2: Add 2 to the result from Step 1: 2 + 12 = 14.
Verification: d ee p hub To verify the answer, I can double-check the
order of operations and the calculations. Multiplication is
indeed performed before addition, and the calculations are correct.
<|special_token|> Summary: The answer is 14.
后处理优化
<think> ummm... multiply 3 and 4... get 12... then add 2...</think> <answer> 14 </answer>
<|special_token|> Reasoning: d ee p hub To solve this, we use order of operations, doing multiplication before addition. Step 1: Multiply 3 by 4, which is 12. Step 2: Add 2 to the result: 2 + 12 = 14. <|special_token|> Summary: The answer is 14.
# 模拟的 R1 Zero 模型粗糙输出 messy_output = "<think> ummm... multiply 3 and 4... get 12... then add 2...</think>\n<answer> 14 </answer>"
def refine_output(messy_text):
think_content = messy_text.split(“<think>”)[1].split(“</think>”)[0].strip()
answer_content = messy_text.split(“<answer>”)[1].split(“</answer>”)[0].strip()refined_text = f"“”<|special_token|> Reasoning: {think_content.replace(‘umm…’, ‘’).strip().capitalize()}.
<|special_token|> Summary: The answer is {answer_content}.“”"
return refined_textrefined_output_text = refine_output(messy_output)
print(“Messy Output (Simulated R1 Zero):”)
print(messy_output)
print(“\nRefined Output:”)
print(refined_output_text)
Messy Output (Simulated R1 Zero): <think> ummm... multiply 3 and 4... get 12... then add 2...</think> <answer> 14 </answer>
Refined Output:
<|special_token|> Reasoning: Multiply 3 and 4… get 12… then add 2.
<|special_token|> Summary: The answer is 14.
在生成冷启动数据之后,下一个关键步骤是 监督微调 (SFT),我们将在下一节详细探讨 SFT 的训练过程。
基于冷启动数据的阶段 1 SFT 训练
为了深入理解 SFT 的训练机制,我们需要了解 SFT 训练器在处理训练数据时,其内部执行了哪些操作?
R1 阶段 1 SFT 训练配置
-
学习清晰的推理风格:以易于理解和追溯的方式组织 “思考” 过程。
-
保持语言一致性:在单个回复中坚持使用单一语言,避免语言混用造成的困扰。
#### 输出 ####
{
'system': 'Your role as an assistant involves ... ',
'conversations': [{'from': 'user', 'value': 'Return your ...'}]
}
#### 输出 ####
# 模型和输出配置 (与之前相同,或根据需要调整) MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct" OUTPUT_DIR = "data/Qwen-SFT-training" # SFT 模型的新输出目录 os.makedirs(OUTPUT_DIR, exist_ok=True)
训练参数 - 与 GRPO 类似,但需针对 SFT 进行调整
training_args = TrainingArguments(
output_dir=OUTPUT_DIR,
overwrite_output_dir=True,
num_train_epochs=1, # 根据需要调整 epoch 数
per_device_train_batch_size=8,
per_device_eval_batch_size=16,
gradient_accumulation_steps=2,
learning_rate=2e-5, # 为 SFT 调整学习率
warmup_ratio=0.1,
weight_decay=0.01,
logging_steps=10,
evaluation_strategy=“no”,
eval_steps=50,
save_strategy=“steps”,
save_steps=50,
save_total_limit=2,
dataloader_num_workers=2,
seed=42,
bf16=True,
push_to_hub=False,
gradient_checkpointing=True,
report_to=“none”,
packing=True, # 启用数据打包以提高训练效率
max_seq_length=4096 # 设置最大序列长度
)模型配置 - 与之前相同
model_args = ModelConfig(
model_name_or_path=MODEL_NAME,
model_revision=“main”,
torch_dtype=“bfloat16”,
trust_remote_code=True,
attn_implementation=“flash_attention_2”
)
阶段 1 SFT 训练循环
# 加载 Bespoke-Stratos-17k 数据集 dataset_sft = load_dataset("HuggingFaceH4/Bespoke-Stratos-17k", split='train') # 仅使用训练集以简化流程
初始化 tokenizer - 与之前相同
tokenizer = AutoTokenizer.from_pretrained(
MODEL_NAME,
trust_remote_code=True,
padding_side=“right”
)
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
# 初始化 SFT 训练的基础模型 - 与之前相同 model_sft = AutoModelForCausalLM.from_pretrained( MODEL_NAME, trust_remote_code=True, torch_dtype=torch.bfloat16 )
初始化 SFT 训练器
sft_trainer = SFTTrainer(
model=model_sft, # 初始化的 Qwen 模型
train_dataset=dataset_sft, # Bespoke-Stratos-17k 数据集
tokenizer=tokenizer, # Tokenizer
args=training_args, # 训练参数
dataset_text_field=“conversations”, # 数据集中包含文本的字段 - SFT 训练的关键参数
packing=True, # 启用数据打包
max_seq_length=4096 # 最大序列长度
)启动 SFT 训练循环
sft_train_result = sft_trainer.train()
... INFO:__main__:Step 10: Loss = ..., Learning Rate = ... INFO:__main__:Step 20: Loss = ..., Learning Rate = ... ...
保存 Tiny R1 LLM
经过阶段 1 SFT 训练微调的模型,即为本文所称的 R1 模型。
面向推理的强化学习
拒绝采样
阶段 2 SFT 训练
-
助人性 (Helpfulness):模型回复是否实用且信息量丰富?
-
无害性 (Harmlessness):模型回复是否安全、无偏见且符合伦理道德?