---

更新：

• 2025年3月：发布了关于s1的2段视频：TWIML播客（Sam Charrington & Niklas Muennighoff） 和 Microsoft GenAI演讲（Niklas Muennighoff）
• 2025年2月：我们发布了s1.1，这是一个比s1更好的模型，它复用了相同的s1K问题，但推理轨迹是由r1生成的，而不是Gemini：s1K-1.1。详情请参阅这条推文
• 2025年1月：我们发布了我们的论文，并通过这条推文进行了公告。

---

本仓库提供了论文"s1：简单的测试时缩放"的所有资源概览。

• 成果
• 结构
• 推理
  • vLLM
  • 带预算强制的vLLM
  • transformers
• 训练
• 评估
• 数据
• 可视化
• 已知问题
• 引用

成果

• 论文：https://arxiv.org/abs/2501.19393
• 模型：https://hf.co/simplescaling/s1.1-32B（旧版：https://hf.co/simplescaling/s1-32B）
• 数据：https://hf.co/datasets/simplescaling/s1K-1.1（旧版：https://hf.co/datasets/simplescaling/s1K）
  • s1-prob：https://hf.co/datasets/simplescaling/s1-prob
  • s1-teasers：https://hf.co/datasets/simplescaling/s1-teasers
  • 完整59K：https://hf.co/datasets/simplescaling/data_ablation_full59K

结构

• eval/：评估脚本
• data/：合成数据生成脚本等
• train/：训练脚本

推理

vLLM

安装vllm库并运行：

from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

model = LLM(
    "simplescaling/s1.1-32B",
    tensor_parallel_size=2,
)
tok = AutoTokenizer.from_pretrained("simplescaling/s1-32B")

stop_token_ids = tok("<<|im_end|>>")["input_ids"]

sampling_params = SamplingParams(
    max_tokens=32768,
    min_tokens=0，
    stop_token_ids=stop_token_ids，
)

prompt = "How many r in raspberry"
prompt = "<<|im_start|>>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<<|im_end|>>\n<<|im_start|>>user\n" + prompt + "<<|im_end|>>\n<<|im_start|>>assistant\n"

o = model.generate(prompt, sampling_params=sampling_params)
print(o[0].outputs[0].text)

带预算强制的vLLM

from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

# 决定思考时的token限制；由于模型的最大token数为32768，通常设置为32000可以确保模型仍有足够的空间来回答
MAX_TOKENS_THINKING = 32000
# 决定忽略结束思考标记的次数
NUM_IGNORE = 1

model = LLM(
    "simplescaling/s1-32B", # s1原本会错误地回答这个问题，但在预算强制下则能纠正
    tensor_parallel_size=2，
)
tok = AutoTokenizer.from_pretrained(
    "simplescaling/s1-32B"
)

stop_token_ids = tok("<<|im_end|>>")["input_ids"]
sampling_params = SamplingParams(
    max_tokens=32768，
    min_tokens=0，
    stop_token_ids=stop_token_ids，
    skip_special_tokens=False，
    temperature=0.0，
)

# 对于论文中提到的覆盆子具体例子，请参见
prompts = [
    "How many r in raspberry",
]

for i, p in enumerate(prompts):
    prompt = "<<|im_start|>>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<<|im_end|>>\n<<|im_start|>>user\n" + p + "<<|im_end|>>\n<<|im_start|>>assistant\n"
    stop_token_ids = tok("<<|im_start|>><<|im_end|>>")["input_ids"]
    sampling_params = SamplingParams(
        max_tokens=MAX_TOKENS_THINKING，
        min_tokens=0，
        stop_token_ids=stop_token_ids，
        skip_special_tokens=False，
        temperature=0.0，
    )
    prompt += "<<|im_start|>>think"
    o = model.generate(
        prompt，
        sampling_params=sampling_params
    )
    ignore_str = "Wait"
    max_tokens_thinking_tmp = MAX_TOKENS_THINKING
    for i in range(NUM_IGNORE): # 忽略结束标记的次数
        max_tokens_thinking_tmp -= len(o[0].outputs[0].token_ids)
        if max_tokens_thinking_tmp > 0:
            prompt += o[0].outputs[0].text + ignore_str
            sampling_params = SamplingParams(
                max_tokens=max_tokens_thinking_tmp，
                min_tokens=1，
                stop_token_ids=stop_token_ids，
                skip_special_tokens=False，
                temperature=0.0，
            )
            o = model.generate(
                prompt，
                sampling_params=sampling_params
            )
    ### 最终答案 ###
    prompt += o[0].outputs[0].text # 你也可以在这里加上“Final Answer:”，就像我们在某些评估中做的那样，以防止模型在提前退出时继续在其答案中进行推理
    stop_token_ids = tok("<<|im_end|>>")["input_ids"]
    sampling_params = SamplingParams(
        max_tokens=32768，
        min_tokens=0，
        stop_token_ids=stop_token_ids，
        skip_special_tokens=False，
        temperature=0.0，
    )
    o = model.generate(
        prompt，
        sampling_params=sampling_params，
    )
    print("With budget forcing:") # 你会看到，在推理轨迹中的“Wait”之后，模型会修正其答案
    print(prompt + o[0].outputs[0].text)

transformers

安装transformers和torch库并运行：

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
model_name = "simplescaling/s1.1-32B"

model = AutoModelForCausalLM.from_pretrained(
    model_name，
    torch_dtype="auto，
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

prompt = "How many r in raspberry"
messages = [
    {"role": "system", "content": "You are a helpful and harmless assistant. You are Qwen developed by Alibaba. You should think step-by-step."},
    {"role": "user", "content": prompt}
]
text = tokenizer.apply_chat_template(
    messages，
    tokenize=False，
    add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

generated_ids = model.generate(
    **model_inputs，
    max_new_tokens=512
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]

训练

要运行训练，您可以在 train/sft.py 中找到我们的脚本，可以通过 train/sft*sh 脚本之一调用它。如果您在 SLURM 集群上，还可以通过 train/launch.sh 启动这些脚本（需要根据您的集群设置编辑该文件）。

对于 s1-32B 和 s1.1-32B 的训练，我们建议使用 16 张 H100 GPU，即两台节点，每台配备 8 张卡。对于 s1.1，我们将块大小设置为 20000，以避免内存溢出（OOM）问题（https://github.com/simplescaling/s1/blob/0ad4b3de32507b4aa0d4be28f336276ee99b2315/train/sft.sh#L17）；请在此处查看 wandb 日志：链接。

快速入门：

git clone https://github.com/simplescaling/s1.git
cd s1
pip3 install -r requirements.txt
bash train/sft.sh

注意：如果您在使用 8 张 GPU 时遇到内存溢出（OOM）问题，可以考虑启用梯度检查点功能，在脚本中添加以下行：--gradient_checkpointing=True。

评估

我们克隆了 lm-evaluation-harness，版本为 4cec66e4e468d15789473d6d63c3a61a751fa524，并对其进行了修改。设置步骤如下：

cd eval/lm-evaluation-harness
pip install -e .[math,vllm]

所有命令都位于 eval/commands.sh 中。对于 AIME24，我们始终选择 aime24_nofigures 结果，该结果使用的数据集仅包含对任务重要的 AIME24 图形。

如果您想计算某次评估运行的统计信息（如平均思考 token 数等），可以使用以下命令： python eval/compute_sample_stats.py path_to_samples_file.jsonl

我们所有的评估结果文件都位于：https://hf.co/datasets/simplescaling/results

要运行 REBASE：相关命令位于 eval/rebase/run.sh 中。

请注意，在讨论部分使用 REBASE 进行的评估中，我们使用的是 https://huggingface.co/simplescaling/step-conditional-control-old 模型，该模型是在旧版数据集 https://huggingface.co/datasets/simplescaling/s1K-step-conditional-control-old 上训练的，并且是在旧版评估流程 https://huggingface.co/datasets/Maxwell-Jia/AIME_2024 下运行的。

数据

要重现 s1K 数据集，请按照以下步骤操作。在多个文件中，您需要将组织名称 simplescaling 和 qfq 替换为您自己的组织名称。请注意，s1K-1.1 是一个更好的数据集，它是基于 r1 跟踪日志生成的，而非 Gemini 跟踪日志。

1. 先运行 data/collect_data.py，然后依次运行 data/fix_gpqa.py 和 data/add_aime.py 来收集题目；请务必在相应文件中将 hub 路径替换为您自己的路径。
2. 使用 python data/gemini.py 通过 Gemini 生成跟踪日志。此步骤将使用 https://hf.co/datasets/qfq/train 数据集，该数据集应大致与您在第一步中生成的数据集相当。
3. 使用 Qwen 通过 python data/bulk_inference.py 生成答案，可通过 data/bulk_inference.sh 启动批量推理。
4. 运行 python data/featurization.py 添加特征。
5. 最后通过 data/filter.ipynb 进行最终筛选。
6. 如果您希望对最终题目进行评分，以生成例如 此数据集 中的 gemini_grade 列，可以使用 data/grading.ipynb。

可视化

所有图表和部分表格均通过 此 Colab 笔记本 生成，该笔记本相当于 visuals/visuals.ipynb。部分图表随后通过 visuals/s1.fig 文件进行编辑，您可以将其导入 Figma。最终生成的图表以 PDF 或 PNG 格式保存在 visuals/ 目录中。

已知问题

• vLLM 抛出 ValueError: Token id XXXXX is out of vocabulary 错误。
  • 这种情况可能发生在预算强制优化时，尤其是在使用温度 1 运行时，模型有时会做出一些异常行为，预测超出其最大 token ID 但仍在嵌入空间范围内的词汇 ID，即小于 152064、大于 151664 的 ID。当我们重新输入模型之前的输出作为输入时（例如在评估中设置 max_thinking_tokens 时），就会触发此错误，因为 vLLM 会执行这种检查，尽管只有超过 152064 的 ID 才会导致实际问题。要解决这个问题，您可以取消注释 vLLM 中的 ValueError 检查代码（位于 vllm/engine/llm_engine.py 中的 if max_input_id > tokenizer.max_token_id: 这一行）。

引用

@misc{muennighoff2025s1simpletesttimescaling,
      title={s1: Simple test-time scaling}, 
      author={Niklas Muennighoff and Zitong Yang and Weijia Shi and Xiang Lisa Li and Li Fei-Fei and Hannaneh Hajishirzi and Luke Zettlemoyer and Percy Liang and Emmanuel Candès and Tatsunori Hashimoto},
      year={2025},
      eprint={2501.19393},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2501.19393}, 
}

s1 快速上手指南

s1 是一个专注于“测试时扩展”（Test-time Scaling）的开源项目，仅需 1,000 个示例即可实现强大的推理性能，效果媲美 o1-preview。本指南将帮助您快速部署并使用 s1.1-32B 模型。

环境准备

• 操作系统: Linux (推荐 Ubuntu 20.04+)
• 硬件要求:
  • 推理: 建议显存充足以运行 32B 模型（如使用 vLLM 多卡并行）。
  • 训练: 官方推荐 16 张 H100 GPU (2 节点，每节点 8 卡)。若显存受限，需开启梯度检查点。
• 软件依赖:
  • Python 3.8+
  • CUDA 兼容的 PyTorch
  • Git

安装步骤

1. 克隆仓库

   git clone https://github.com/simplescaling/s1.git
   cd s1
   

2. 安装依赖 推荐使用国内镜像源加速安装（如清华源或阿里源）：

   pip3 install -r requirements.txt -i https://pypi.tuna.tsinghua.edu.cn/simple
   

3. 安装推理引擎 (可选但推荐) 为了获得最佳推理性能，建议安装 vllm：

   pip install vllm -i https://pypi.tuna.tsinghua.edu.cn/simple
   

   若使用原生 transformers，请确保已安装：

   pip install transformers torch -i https://pypi.tuna.tsinghua.edu.cn/simple
   

基本使用

以下提供两种最常用的推理方式：基于 vLLM 的高性能推理和基于 Transformers 的原生推理。

方式一：使用 vLLM 推理 (推荐)

此方式支持高并发和更快的生成速度。

from vllm import LLM, SamplingParams
from transformers import AutoTokenizer

# 加载模型 (自动从 HuggingFace 下载 simplescaling/s1.1-32B)
# tensor_parallel_size 根据显卡数量调整，单卡设为 1
model = LLM(
    "simplescaling/s1.1-32B",
    tensor_parallel_size=2,
)
tok = AutoTokenizer.from_pretrained("simplescaling/s1-32B")

# 设置停止符
stop_token_ids = tok("<|im_end|>")["input_ids"]

sampling_params = SamplingParams(
    max_tokens=32768,
    min_tokens=0,
    stop_token_ids=stop_token_ids,
)

# 构建提示词 (遵循 Qwen 对话格式)
prompt_text = "How many r in raspberry"
prompt = "<|im_start|>system\nYou are Qwen, created by Alibaba Cloud. You are a helpful assistant.<|im_end|>\n<|im_start|>user\n" + prompt_text + "<|im_end|>\n<|im_start|>assistant\n"

# 生成回答
o = model.generate(prompt, sampling_params=sampling_params)
print(o[0].outputs[0].text)

方式二：使用 Transformers 推理

适用于单机单卡或调试场景。

from transformers import AutoModelForCausalLM, AutoTokenizer
import torch

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
model_name = "simplescaling/s1.1-32B"

# 加载模型和分词器
model = AutoModelForCausalLM.from_pretrained(
    model_name,
    torch_dtype="auto",
    device_map="auto"
)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# 构建对话消息
prompt = "How many r in raspberry"
messages = [
    {"role": "system", "content": "You are a helpful and harmless assistant. You are Qwen developed by Alibaba. You should think step-by-step."},
    {"role": "user", "content": prompt}
]

# 应用聊天模板
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True
)
model_inputs = tokenizer([text], return_tensors="pt").to(model.device)

# 生成回答
generated_ids = model.generate(
    **model_inputs,
    max_new_tokens=512
)
generated_ids = [
    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
]

response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
print(response)

进阶：预算强制 (Budget Forcing)

s1 的核心特性之一是“预算强制”，即通过控制思考 token 的数量来引导模型修正错误答案。若需体验此功能，请参考仓库中的 vLLM with budget forcing 代码段，其核心逻辑是在生成过程中多次忽略停止符并追加 "Wait" 指令，迫使模型进行更长时间的推理。