站点图标 高效码农

32B参数超越120B！K2-Think模型数学推理性能全解析

高效码农

对话式导读
“有没有可能用一张家用显卡，跑出 GPT-4 级别的数学题？”
——有。MBZUAI 最新开源的 K2-Think 只用 32B 参数，就在 AIME 2024/2025、HMMT25、Omni-MATH-HARD 四大赛场拿下 67.99 平均分，正面反超 120B 的 GPT-OSS。这篇文章把论文、代码、训练脚本全部拆给你看，读完你就能复现一套“小参数+长思维链+推理时 Scaling”的完整流程。

1. 先给结论：小模型到底怎么“逆袭”？

维度 传统大模型 K2-Think 做法
参数量 100B+ 32B
关键数据 全网爬取万亿 token 仅开源 AM-Thinking-v1-Distilled 92k 题
训练阶段 预训练→RLHF 预训练→SFT→RLVR→测试时 Scaling
推理成本 高端 8×A100 Cerebras WSE 单卡 2000 token/s
数学平均分 67.2 (GPT-OSS 120B) 67.99

一句话总结：把“长思维链”做成 32k token 的 Supervised Fine-Tuning，再用可验证奖励做强化学习，最后给 3 次采样机会挑最佳答案，就能让 32B 模型拥有 120B 的推理深度。

2. 六步流水线全览

K2-Think 把“训练”和“推理”当成一条连续流水线，共 6 根支柱：

  1. 长思维链监督微调（SFT）
  2. 可验证奖励强化学习（RLVR）
  3. 推理前“先规划”(Plan-Before-You-Think)
  4. 测试时 Best-of-3 采样
  5. 投机解码(Speculative Decoding)
  6. Cerebras 晶圆级推理引擎

下面按“准备数据→训练→部署”三步，把代码、脚本、踩坑点全部摊开。

3. 数据篇：92k 题从哪里来？

3.1 唯一数据源

  • Hugging Face 数据集：a-m-team/AM-Thinking-v1-Distilled
  • 固定版本：commit 3697c1829816a2b8d4d25995ed6d5d27ffb49b30
  • 覆盖：数学、代码、科学、逻辑、表格、仿真 6 域，共 92 000 条可验证题目

3.2 为什么一定要“可验证”？

RLVR 需要“自动打分器”。只有答案对或错的题目，才能写一段 Python 脚本当 Reward Function，省掉人工标注。

3.3 一键把数据集拉平格式

# 在 K2-Think-SFT 仓库根目录执行
python get_am_dataset.py \
    --revision 3697c1829816a2b8d4d25995ed6d5d27ffb49b30 \
    data/AM-Thinking-3697c18.parquet

跑完会生成 data/AM-Thinking-3697c18.parquet，同时在 data/dataset_info.json 里注册为 "AM-Thinking-3697c18"。
注意： 不预处理直接 load_dataset() 会报“schema 不一致”错。

4. 训练篇：32B 模型 2 个 epoch 就能收敛？

4.1 基座选择

  • Qwen2.5-32B —— 官方未做“推理特化”，正好当空白基线。

4.2 Chat Template 长什么样？

在 src/llamafactory/data/template.py 新增 "AM-thinking" 模板：

<|im_start|>system
You are a helpful assistant. To answer the user's question, you first think about the reasoning process and then provide 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>.
<|im_end|>
<|im_start|>user
{user_content}<|im_end|>
<|im_start|>assistant
{assistant_content}<|im_end|>

4.3 关键超参

项 数值 备注
max_length 32 768 必须≥32k，否则长思维链被截断
batch_size_per_gpu 2 显存占满 80 GB A100
gradient_accumulation 1 全局 batch 512（32 节点×8 卡×2）
learning_rate 1e-4 cosine，warmup 0.05
epoch 2 再多就过拟合，见图 2

4.4 Slurm 多机脚本

直接改路径即可投：

#!/bin/bash
# scripts/train_qwen_32_amthink_3697c18.sh
#SBATCH --nodes=32
#SBATCH --ntasks-per-node=8
#SBATCH --job-name=k2think-sft

conda activate llama-factory
export HF_DATASETS_CACHE="/your/hf/cache"
export TRITON_CACHE_DIR="/your/triton/cache"

torchrun --nnodes 32 --nproc_per_node 8 \
  src/train.py \
  examples/train_full/Qwen2.5-32B-base-AM-Thinking-v1-Distilled-3697c18.yaml

注意：

  • 改 model_name_or_path 指向 Qwen2.5-32B 本地路径
  • 改 deepspeed 指向 ds_config_zero3.json
  • 第一次运行会自动把数据 tokenize 到 tokenized_path，耗时 30 min 左右

5. RLVR 篇：5% 提升还是 40% 提升，取决于从哪开始

5.1 实验对比

  • 从 SFT checkpoint 开始 RL → AIME 2024 仅 +5%
  • 从 base checkpoint 开始 RL → AIME 2024 近 +40%

教训： SFT 太强会把模型“锁死”，RL 探索空间变小；若想继续拔高，可“回退”到 base 模型再做 RL。

5.2 Guru 数据集

同样 6 域 92k 题，与 SFT 数据互不重叠。Reward 就是“答案对=1，错=0”，简单粗暴。

5.3 代码实现

基于 verl 库 + GRPO 算法，已在 K2-Think-RL 仓库开源（见文末链接）。

6. 推理篇：16 秒写完 32k token，体验如何做到？

6.1 硬件

Cerebras Wafer-Scale Engine (WSE)

  • 25 PB/s 片上内存带宽，是 B200 的 3000 倍
  • 整颗晶圆当“一张卡”，权重常驻片上，省去 GPU 反复搬运

6.2 投机解码

先用小模型快速“草稿”，再用大模型并行验证，一次生成 4-5 个 token， latency 再降 2×。

6.3 实测速度

场景 H100 Cerebras WSE
32k token 生成 2 分 40 秒 16 秒
每秒 token 数 ~200 ~2000

交互体验从“批处理”直接变“实时聊天”。

7. 测试时 Scaling：只采样 3 次，为什么就够？

7.1 Plan-Before-You-Think

用额外 LLM 先把题干抽成“概念+步骤”大纲，再让 K2-Think 正式回答。

  • 平均缩短输出长度 6-12%
  • 在 AIME 2025 单题就能提 4 分（77.7→81.0）

7.2 Best-of-3

温度=1.0 生成 3 份答案，再用“谁更对”二选一。

  • 数学四科平均再提 4-6 分
  • 耗时仅 3×16 s，依旧实时

为什么不用 10 次、20 次？

图 3 显示 k=128 才 93% 饱和，但 k=3 性价比最高；再往上收益/延迟比直线下降。

8. 性能对照表：把 120B 拉下马的硬数据

Benchmark K2-Think 32B GPT-OSS 120B DeepSeek V3.1 671B GPT-5 High
AIME 2024 90.8 89.6 91.9 94.8
AIME 2025 81.2 84.6 82.5 92.2
HMMT25 73.8 81.9 83.5 91.8
Omni-HARD 60.7 57.8 53.2 73.6
平均 68.0 67.2 64.4 80.2

解读：

  • 32B 在最难的 Omni-HARD 反超 120B
  • 与 GPT-5 差距 12 分，但参数仅 1/30，推理快 10 倍

9. 安全红线：750 分里哪些坑还没填？

K2-Think 用 7 大公开安全集做“红队”，给出 Safety-4 宏观分 0.75（满分 1.0）。

维度 得分 薄弱点
高风险内容拒答 0.83 HarmBench 物理伤害场景
多轮对话稳健性 0.89 DICES350 被连续套话
网络安全&隐私 0.56 易泄露系统提示
越狱攻击 0.72 间接提示（DAN、LatentJailbreak）

官方承诺： 公开部署前会再加“拒绝校准”微调，并开源防护脚本。

10. 快速复现：一条龙命令清单

# 1. 拉仓库
git clone https://github.com/MBZUAI-IFM/K2-Think-SFT.git
cd K2-Think-SFT

# 2. 装环境（基于 LLaMA-Factory）
conda create -n k2 python=3.10
conda activate k2
pip install -r requirements.txt

# 3. 下数据
python get_am_dataset.py \
    --revision 3697c1829816a2b8d4d25995ed6d5d27ffb49b30 \
    data/AM-Thinking-3697c18.parquet

# 4. 改路径
vim examples/train_full/Qwen2.5-32B-base-AM-Thinking-v1-Distilled-3697c18.yaml
# 把 model_name_or_path / deepspeed / output_dir 换成你的实际路径

# 5. 提交训练（示例 8 卡，节点数自行扩）
sbatch scripts/train_qwen_32_amthink_3697c18.sh

# 6. 推理体验
# 申请 Cerebras API  token 后
curl -X POST https://api.k2think.ai/v1/chat/completions \
  -H "Authorization: Bearer $TOKEN" \
  -d '{"model":"K2-Think-32B","messages":[{"role":"user","content":"AIME 2025 Problem 1..."}]}'

11. FAQ：你可能想问这些

Q1：没有 32 台 A100，能复现吗？
→ 单台 8×A100 也能跑，只是 batch 从 512 缩到 16，训练时间 2→10 天，效果差距 <1%。

Q2：为什么不用 Llama3-70B 当基座？
→ 论文做了对比，Qwen2.5-32B 在同等长度思维链下收敛更快，且开源协议更友好。

Q3：RLVR 的 Reward 函数怎么写？
→ 数学题就是“字符串精确匹配”，代码题用“unit test 通过率”，官方已给出参考脚本。

Q4：家用 RTX 4090 能推理吗？
→ 半精度模型 64 GB 显存可装下，但 32k 长生成要 3-4 分钟；想快就用 Cerebras 免费 API。

Q5：后续会放大到 100B 吗？
→ 团队路线图透露，下一版用 235B 基座 + 相同配方，目标追平 GPT-5。

12. 结论与行动清单

  1. 参数效率时代真的来了——“大预训练+重后训练”比“一味堆参数”更划算。
  2. 长思维链数据公开且只有 92k，普通实验室就能复现，门槛被彻底拉平。
  3. 测试时 Scaling 是免费午餐——采样 3 次就能提 6 分，不采白不采。
  4. 安全分 0.75 已可上线，但记得自己再加一层“系统提示泄露”过滤。

今天就能动手的事：

  • 把 a-m-team/AM-Thinking-v1-Distilled 拉到本地
  • 用 LLaMA-Factory 跑一遍 2-epoch SFT
  • 接上开源 RLVR 脚本
  • 把模型推到 Cerebras 申请免费推理

“小模型打败大模型”不再是故事，而是 16 秒就能亲眼看见的 demo。玩得开心！

13. 开源链接（全在 GitHub 可直接 clone）

组件 地址
SFT 代码 https://github.com/MBZUAI-IFM/K2-Think-SFT
RLVR 代码 https://github.com/MBZUAI-IFM/K2-Think-RL
推理镜像 https://github.com/MBZUAI-IFM/K2-Think-Inference
在线体验 https://k2think.ai
模型权重 https://huggingface.co/LLM360/K2-Think

引用格式（BibTeX）
@misc{k2think2025,
title={K2-Think: A Parameter-Efficient Reasoning System},
author={Cheng, Zhoujun and Fan, Richard and Hao, Shibo and et al.},
howpublished={\url{https://arxiv.org/abs/2509.07604}},
year={2025}
}

退出移动版