AI实战 | SFT微调1.5B数学模型,喂R1推理路径能提多少?

从数据清洗到训练到评测，完整代码实战

🎯 问题：大模型做数学推理靠的是"想得久"——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评测

01.

环境准备

核心依赖就这些：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

02.

数据集加载

使用DeepMath-103K数学推理训练集。这数据集几个特点：103K个高难度数学样本，每个带难度标注和来自DeepSeek-R1的三条推理路径，而且对主流数学测试集做了去重，训练集里不会出现测试集的题。

# 使用datasets的api来加载数据集
data = load_dataset("zwhe99/DeepMath-103K")

看一下数据结构：

data

DatasetDict({

train: Dataset({

features: ['question', 'final_answer', 'difficulty', 'topic', 'r1_solution_1', 'r1_solution_2', 'r1_solution_3'],

num_rows: 103022

})

03.

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|>

04.

数据预处理

两步预处理：第一步，筛选答案是纯数字的题目（方便后续评测）；第二步，按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小模型的硬限制。

05.

模型训练

用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")

06.

推理评测

用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.0

from 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%准确率

07.

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内容，关注主页查看~