从数据清洗到训练到评测,完整代码实战
🎯 问题:大模型做数学推理靠的是"想得久"——DeepSeek-R1的推理路径动辄几千token,自我反思、思路切换、反复验算。那问题来了:把R1的推理路径喂给一个1.5B的小模型做SFT,它能学到什么程度?
这篇文章是完整的代码实战。模型用Qwen2.5-Math-1.5B,数据用DeepMath-103K,评测用MATH500。从环境搭建、数据加载预处理、模型训练,到vLLM推理和数学评测,全部代码贴出来,直接能跑。
🔧 方案:SFT全流程代码实战——DeepMath-103K数据 + Qwen2.5-Math-1.5B基座 + HuggingFace Trainer微调 + vLLM推理 + MATH500评测
环境准备
核心依赖就这些:transformers做训练,datasets加载数据,vllm做推理。实验用的是8卡A100环境,单卡也能跑(batch size调小就行)。
import transformers
import torch
import datasets
from datasets import load_dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForSeq2Seq
import matplotlib.pyplot as plt
from vllm import LLM, SamplingParams
import copy看一眼版本,确认环境没问题:
print("Transformers version:", transformers.__version__)
print("Torch version:", torch.__version__)
print("Datasets version:", datasets.__version__)
print("VLLM version:", transformers.__version__)输出:
Transformers version: 4.52.3
Torch version: 2.7.0+cu126
Datasets version: 3.6.0
VLLM version: 4.52.3
print("Number of CUDA devices:", torch.cuda.device_count())
print("Current CUDA device:", torch.cuda.current_device())Number of CUDA devices: 8
Current CUDA device: 0
数据集加载
使用DeepMath-103K数学推理训练集。这数据集几个特点:103K个高难度数学样本,每个带难度标注和来自DeepSeek-R1的三条推理路径,而且对主流数学测试集做了去重,训练集里不会出现测试集的题。
# 使用datasets的api来加载数据集
data = load_dataset("zwhe99/DeepMath-103K")看一下数据结构:
dataDatasetDict({
train: Dataset({
features: ['question', 'final_answer', 'difficulty', 'topic', 'r1_solution_1', 'r1_solution_2', 'r1_solution_3'],
num_rows: 103022
})
})
Tokenizer与chat_template
Tokenizer负责把自然语言转成模型能接受的整数id列表。chat_template是大模型特有的东西,定义了怎么把对话组织成特定格式。Qwen2.5-Math的默认system prompt会告诉模型"逐步推理,最终答案放在\boxed{}里"。
# 加载Tokenizer
qwen_tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Math-1.5B")
print(qwen_tokenizer.chat_template)用apply_chat_template把对话转成带特殊token的格式,感受一下实际输入长什么样:
# 通过apply_chat_template方法可以把一个对话历史转变成带有特殊token的输入格式
# tokenize=False表示不进行分词,add_generation_prompt=False表示不添加生成提示
text = "hello world"
print(qwen_tokenizer.apply_chat_template(
[{"role": "user", "content": text}, {"role": "assistant", "content": "hi!"}],
tokenize=False, add_generation_prompt=False
))<|im_start|>system
Please reason step by step, and put your final answer within \boxed{}.<|im_end|>
<|im_start|>user
hello world<|im_end|>
<|im_start|>assistant
hi!<|im_end|>
数据预处理
两步预处理:第一步,筛选答案是纯数字的题目(方便后续评测);第二步,按token长度筛选,只保留4096以内的样本。
筛选纯数字答案
# 数据筛选:只保留答案是纯数字的
data_with_number_answer = data["train"].filter(
lambda x: x["final_answer"].isdigit()
)
print(len(data_with_number_answer))Tokenize:把问题+R1推理路径转成训练格式
每个问题取第一条R1推理路径作为标准答案,用chat_template组织成对话格式,然后在assistant回复前加上<think>标记——这是Qwen2.5-Math的思考模式触发词。
# 数据预处理:每个问题以第一个DeepSeek-R1的输出作为标准答案
# 将其处理成包含长度、input_ids和labels的格式
def tokenize(example):
question = example["question"]
r1_solution = example["r1_solution_1"]
message = [
{"role": "user", "content": question},
{"role": "assistant", "content": f"<think>\n{r1_solution}\n"}
]
result = qwen_tokenizer.apply_chat_template(
message, add_generation_prompt=False,
tokenize=True, return_dict=True
)
result["length"] = len(result["input_ids"])
result["labels"] = result["input_ids"].copy()
return result
train = data_with_number_answer.map(tokenize, batched=False, num_proc=16)解码一个样本看看实际训练输入长什么样——R1的回复包含大量自我反思、思路切换和反复验算:
# 数据示例:R1的回复中包含自我反思、思路切换、验算等高级推理特征
print(qwen_tokenizer.decode(train[0]["input_ids"], skip_special_tokens=False))输出示例(截取开头):
<|im_start|>system
Please reason step by step, and put your final answer within \boxed{}.<|im_end|>
<|im_start|>user
Evaluate the limit: [极限题]<|im_end|>
<|im_start|>assistant
<think>
Okay, so I have this limit to evaluate...
...[大量推理过程]...
\boxed{0}<|im_end|>
长度筛选
先看长度分布——R1推理路径的平均长度在5000左右,但Qwen2.5-Math-1.5B的最大长度只有4096,所以只保留4096以内的样本。
# 长度分布可视化:R1的回复的长度均值约为5000
lengths = train["length"]
plt.hist(lengths, bins=100)
plt.xlabel("Length")
plt.ylabel("Count")
plt.title("Length Distribution")
plt.show()# 数据筛选:只保留长度小于4096的样本
train = train.filter(lambda x: x["length"] < 4096, num_proc=16)
print(len(train))15310
103K→筛纯数字→筛长度,最终留下约15K条训练样本。R1推理路径太长,大量样本超4096被砍掉了——这是1.5B小模型的硬限制。
模型训练
用Qwen2.5-Math-1.5B做基座。注意这是个基座模型,没经过指令微调,直接用只能做ICL生成。用BF16精度加载,单张80GB A100就能跑。
加载模型
# 加载模型
model = AutoModelForCausalLM.from_pretrained(
"Qwen/Qwen2.5-Math-1.5B",
torch_dtype=torch.bfloat16,
device_map="auto"
)
model模型结构(28层Qwen2DecoderLayer,1536隐藏维度):
Qwen2ForCausalLM(
(model): Qwen2Model(
(embed_tokens): Embedding(151936, 1536)
(layers): 28 x Qwen2DecoderLayer(...)
)
(lm_head): Linear(1536, 151936, bias=False)
)
训练参数
# 设置训练的参数
training_args = TrainingArguments(
# batch size & epochs
per_device_train_batch_size=4, # 显存不足时,请设置为1
gradient_accumulation_steps=16,
num_train_epochs=3,
# hyperparameters
learning_rate=1e-6,
lr_scheduler_type="cosine",
# monitoring
output_dir="./checkpoints",
logging_dir="./checkpoints/logs",
report_to="tensorboard",
eval_strategy="no",
save_strategy="steps",
save_steps=100,
save_total_limit=1,
# efficiency
bf16=True,
group_by_length=True,
torch_compile=True,
gradient_checkpointing=False, # 显存不足时,请设置为True
# reproducibility
seed=42,
data_seed=42
)
# 设置默认的用于seq2seq任务的DataCollator
collator = DataCollatorForSeq2Seq(
qwen_tokenizer,
padding=True,
pad_to_multiple_of=8,
return_tensors="pt",
)启动训练
# 使用Trainer进行训练
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train,
processing_class=qwen_tokenizer,
data_collator=collator,
)
trainer.train()# 训练完毕后,保存最终模型
trainer.save_model("./checkpoints/final_model")推理评测
用MATH500测试集评测。MATH500是MATH测试集的子集,500道高中竞赛级别数学题,每道都有难度标注。
# 加载数据
math500 = load_dataset("zwhe99/MATH", split="math500")看一条数据长什么样:
# 观察数据格式
math500[0]输出:
{'problem': 'Convert the point (0,3)...', 'level': 'Level 2', 'type': 'Precalculus', 'expected_answer': '(3,\\frac{\\pi}{2})'}
预处理测试数据
## 预处理:将问题处理成适合模型输入的对话格式
data = []
for line in math500:
line = copy.deepcopy(line)
line["messages"] = [{"role": "user", "content": line["problem"]}]
data.append(line)
print(len(data))500
vLLM推理
用vLLM做推理后端,速度快。base_model填训练完保存的模型路径。
base_model = "./checkpoints/final_model"
llm = LLM(
model=base_model,
tensor_parallel_size=1,
max_model_len=4096
)
tokenizer = AutoTokenizer.from_pretrained(base_model)# 将输入转变成Prompt的格式:VLLM直接推理需要接受经过apply_chat_template方法处理的输入
prompts = [
tokenizer.apply_chat_template(
conversation=line["messages"],
tokenize=False,
add_generation_prompt=True,
)
for line in data
]
# 和训练格式一致:添加<think>标记激发模型思考
prompts = [prompt + "<think>\n" for prompt in prompts]
# 生成参数设置:temperature=0表示贪婪解码,max_tokens=3700是为了适应模型的最大长度
sampling_params = SamplingParams(temperature=0, max_tokens=3700)
# 实际生成
outputs = llm.generate(prompts, sampling_params)数学评测:抽取答案+比对
用规则匹配判断答案是否正确。模型生成的答案一般包在\boxed{}里,用正则提取出来再和标准答案比对。
# 安装必要的依赖
!pip install antlr4-python3-runtime==4.11.0from openmathinst_utils import math_equal, extract_answer# 处理生成的输出,将答案提取出来
for line, output in zip(data, outputs):
line["response"] = output.outputs[0].text
line["extracted_answer"] = extract_answer(output.outputs[0].text)# 使用math_equal方法来验证提取的答案是否正确
for line in data:
line["correct"] = math_equal(line["extracted_answer"], line["expected_answer"])
print("Accuracy:", sum(line["correct"] for line in data) / len(data))✅ 效果
1.5B模型 + 15K条SFT数据 + 3个epoch,在MATH500上达到45.2%准确率
Case Study:看看模型学会了什么
随机看一条模型输出:
import random
item = random.choice(data)
print(item["problem"])
print("="*50)
print(item["response"])两个观察:1) 微调后的模型有了基本的指令遵循能力,能按步骤解题;2) 模型出现明显的反思倾向,倾向于用代码验证答案——这些特征直接来自R1的推理路径。
💥 踩坑
长度截断丢数据:R1推理路径平均5000 tokens,1.5B模型最大长度4096,大量样本直接被砍掉。103K→15K,损失了85%的数据。
labels复制了input_ids:代码里labels=input_ids.copy(),模型在学整个序列(包括system prompt和用户问题),不只是assistant部分。正经做法应该把非assistant部分的labels设为-100。
学习率1e-6偏保守:1.5B模型用1e-6的lr做SFT收敛很慢。通常SFT的学习率在1e-5到5e-5之间。
只用了第一条R1推理路径:DeepMath-103K每条数据有3条R1推理路径,只用了第一条,浪费了2/3的监督信号。
💡 结论
SFT确实能让1.5B小模型学到R1的推理"姿势"——微调后模型有了指令遵循能力,会出现自我反思和代码验算的行为。45.2%的准确率说明R1的推理模式能被蒸馏下来,但受限于模型容量和序列长度,离R1本身的水平还差很远。
关键瓶颈不是训练方法,而是序列长度——1.5B模型的4096窗口塞不下R1动辄上万的推理过程。想进一步提升,要么用支持长序列的模型,要么截断推理路径只保留核心步骤,要么走GRPO之类的强化学习路线让模型自己学会"精简推理"。
📌 更多AI内容,关注主页查看~
夜雨聆风