小米MiMo系列:MiMo-7B-RL技术报告解读
文章目录
- 1、模型简介
- 2、相关背景
- 3、主要创新点
- 预训练:为推理而生的基础模型
- 后训练:开创性的推理模型
- RL 基础设施
- 4、模型细节
- 5、评测结果
- 6、模型部署
- SGLang 推理
- vLLM 推理
- HuggingFace 推理
项目开源
- 技术报告:https://arxiv.org/abs/2505.07608
- github:https://github.com/XiaomiMiMo/MiMo
1、模型简介
MiMo-7B系列模型是从零开始训练并专为推理任务而生的。其强化学习实验显示,模型具有非凡的推理潜力,甚至超过了大得多的320亿参数模型。此外,在冷启动的SFT模型上进行了RL训练,生成了MiMo-7B-RL,该模型在数学和代码推理任务上表现出色,与OpenAI o1-mini的表现相当。
MiMo-7B系列包括:基础模型(MiMo-7B-Base)、SFT模型(MiMo-7B-SFT)、从基础模型训练的RL模型(MiMo-7B-RL-Zero)以及从SFT模型训练的RL模型(MiMo-7B-RL)。
2、相关背景
目前,大多数成功的强化学习工作,都依赖于相对较大的基础模型,例如320亿参数的模型,特别是在增强代码推理能力方面。此外,普遍认为在一个小模型中同时提高数学和代码能力是具有挑战性的。
然而,我们认为强化学习训练的推理模型的有效性依赖于基础模型固有的推理潜力。为了充分释放语言模型的推理潜力,不仅需要关注后训练,还需要关注针对推理定制的预训练策略。
3、主要创新点
预训练:为推理而生的基础模型
- 优化了数据预处理流程,增强了文本提取工具包,并应用了多维度数据过滤,以增加预训练数据中的推理模式密度。还采用了多种策略来生成大量多样化的合成推理数据。
- 采用三阶段数据混合策略进行预训练。总体而言,MiMo-7B-Base 在大约 25 万亿个令牌上进行了预训练。
- 将多令牌预测MTP作为额外的训练目标纳入其中,这增强了模型性能并加速了推理。
后训练:开创性的推理模型
- 精心挑选了 130K 个数学和代码问题作为强化学习(RL)训练数据,这些问题可以通过基于规则的验证器进行验证。每个问题都经过仔细清理和难度评估,以确保质量。仅使用基于规则的准确性奖励,以避免潜在的奖励黑客攻击。
- 为了缓解具有挑战性的代码问题的稀疏奖励问题,引入了测试难度驱动的代码奖励。通过对不同难度级别的测试用例分配细粒度的分数,可以通过密集的奖励信号更有效地优化策略。
- 对简单问题实施了数据重采样策略,以提高 rollout 采样效率并在 RL 训练的后期阶段稳定策略更新。
RL 基础设施
- 开发了一个无缝 Rollout 引擎,以加速 RL 训练和验证。设计集成了连续 rollout、异步奖励计算和提前终止,以最小化 GPU 空闲时间,实现了 2.29× 的更快训练和 1.96× 的更快验证。
- 在 vLLM 中支持 MTP,并增强了 RL 系统中推理引擎的鲁棒性。
4、模型细节
MiMo-7B 的 MTP 层在预训练和 SFT 期间进行微调,并在 RL 期间冻结。对于推测解码,一个 MTP 层的接受率约为 90%。
MiMo-7B系列主要包含如下四款模型,其具体描述和下载地址如下:
- Huggingface下载地址: https://huggingface.co/XiaomiMiMo
- modelscope下载地址: https://www.modelscope.cn/organization/XiaomiMiMo
5、评测结果
模型通用能力:基本与QwQ-32B-Preview相当
数学和代码能力:优于OpenAI o1-mini等其他模型
MiMo-7B系列模型对比:其中 MiMo-7B-RL 在 数学和代码能力上最强
注:评估是在 temperature=0.6 的条件下进行的。AIME24 和 AIME25 的分数是 32 次重复的平均分。
6、模型部署
SGLang 推理
- SGLang 支持了 MiMo 模型 和 MTP,示例脚本:
# Install the latest SGlang from main branch
python3 -m uv pip install "sglang[all] @ git+https://github.com/sgl-project/sglang.git/@main#egg=sglang&subdirectory=python"# Launch SGLang Server
SGLANG_USE_MODELSCOPE=true python3 -m sglang.launch_server --model-path XiaomiMiMo/MiMo-7B-RL --host 0.0.0.0 --trust-remote-code# Launch MTP Server
SGLANG_USE_MODELSCOPE=true python3 -m sglang.launch_server --model-path XiaomiMiMo/MiMo-7B-RL --trust-remote-code \
--speculative-algorithm EAGLE --speculative-num-steps 1 --speculative-eagle-topk 1 \
--speculative-num-draft-tokens 2 --mem-fraction 0.5
详细用法请参阅: SGLang 文档
vLLM 推理
官方支持使用 MiMo-vLLM 的分支 进行 MiMo-MTP 推理。示例代码:(推荐使用 vLLM 0.7.3 版本)
from vllm import LLM, SamplingParamsmodel_path = "/path/to/MiMo"
llm = LLM(model=model_path,trust_remote_code=True,num_speculative_tokens=1,disable_log_stats=False
)
sampling_params = SamplingParams(temperature=0.6)conversation = [{"role": "system","content": ""},{"role": "user","content": "Write an essay about the importance of higher education.",},
]outputs = llm.chat(conversation,sampling_params=sampling_params,use_tqdm=False)for output in outputs:prompt = output.promptgenerated_text = output.outputs[0].textprint(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")print("=" * 80)
你也可以注册一个不加载 MTP 参数的 vLLM 加载器来使用 MiMo。你需要将 registry/register_mimo_in_vllm.py 复制到你的目录并通过以下方式导入:
import register_mimo_in_vllmfrom vllm import LLM, SamplingParamsmodel_path = "/path/to/MiMo"
llm = LLM(model=model_path,trust_remote_code=True,# num_speculative_tokens=1,disable_log_stats=False
)
sampling_params = SamplingParams(temperature=0.6)
HuggingFace 推理
示例代码:
from modelscope import AutoModel, AutoModelForCausalLM, AutoTokenizermodel_id = "XiaomiMiMo/MiMo-7B-RL"
model = AutoModelForCausalLM.from_pretrained(model_id, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(model_id)
inputs = tokenizer(["Today is"], return_tensors='pt')
output = model.generate(**inputs, max_new_tokens = 100)
print(tokenizer.decode(output.tolist()[0]))