JustRL-II基础模型

JustRL-II-base-model

二〇二六年九月七日 · 英文原文

摘要

OpenBMB 发布 JustRL-II 基础模型检查点，作为博客《JustRL II：借助 Critic 将小型 LLM 扩展到 128K 推理》中数学推理案例研究的 RL 初始化起点。该检查点为已训练的小型语言模型，支持长思维链输出，在 AIME 2025 上得分约 61%。完整 JustRL II 配方经约 300 步 RL 后达 81%，标准 GRPO 基线停滞于约 74%。训练数据经三阶段 pipeline 审计约 10 万道开源数学题，最终保留 32,412 道。模型采用 Llama 架构，支持 vLLM 或 SGLang 服务。

JustRL-II 基础模型

这是博客 JustRL II：借助 Critic 将小型 LLM 扩展到 128K 推理（中文版）中使用的 RL 初始化检查点。

它是该博客中数学推理案例研究的起点 —— 博客中的每次运行（标准 GRPO 基线、完整的 JustRL II 配方以及所有消融实验）都从该检查点初始化并以其为评估基准。它不是 RL 后的模型。我们发布它是为了让博客的数据 pipeline 和训练配方可以从完全相同的起始权重复现。

该检查点是什么

一个已经过训练、能够生成长思维链（thinking）回复的小型语言模型。它在最终答案之前于 <think> ... </think> 内输出推理过程。
在 RL 之前，按照博客的评估协议（对采样回复取平均，128k token 生成预算），它在 AIME 2025 上得分约为 61%。
从该检查点出发，完整的 JustRL II 配方在约 300 步 RL 后在 AIME 2025 上达到 81%，而相同数据上的标准 GRPO 基线则停滞在约 74%（见博客图 1、2 和 8）。
博客的难度重新校准步骤（数据 pipeline 的第三阶段）通过对该检查点每个问题 rollout 8 次来计算，因此发布的训练集的难度分级对应的是该模型的通过率。

文件

文件	说明
`pytorch_model.bin`	bf16 权重，单文件
`config.json`	Llama 风格架构（`LlamaForCausalLM`），可用标准 `transformers` 加载
`tokenizer.json`、`tokenizer_config.json`、`special_tokens_map.json`	tokenizer
`chat_template.jinja`	支持思考模式（`enable_thinking`）的聊天模板

配置了两个 end-of-sequence id（eos_token_id = [1, 130073]）；在生成调用或 serving 引擎中请同时传入。config.json 自带 max_position_embeddings = 65536；博客中的 RL 运行使用 128k token 生成预算 —— 长上下文 serving 设置请参阅博客。

使用方法

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_id = "openbmb/JustRL-II-base-model"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id, torch_dtype=torch.bfloat16, device_map="auto"
)

messages = [{"role": "user", "content": "What is the sum of all positive divisors of 360? Think step by step."}]
text = tokenizer.apply_chat_template(
    messages, tokenize=False, add_generation_prompt=True, enable_thinking=True
)
inputs = tokenizer(text, return_tensors="pt").to(model.device)
inputs.pop("token_type_ids", None)

out = model.generate(
    **inputs,
    max_new_tokens=8192,
    do_sample=True,
    temperature=1.0,
    eos_token_id=[tokenizer.eos_token_id, 130073],
)
print(tokenizer.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=False))

对于评估或 RL rollout，可将其作为标准 Llama 架构模型用 vLLM 或 SGLang 提供服务，例如：

vllm serve openbmb/JustRL-II-base-model --dtype bfloat16

预期用途

从相同初始化复现 JustRL II 配方及其消融实验。
小型模型长 CoT RL 研究：credit assignment、critic 诊断、数据难度校准。

该检查点未针对通用助手用途进行对齐，且仅经过数学推理评估。它可能产生非常长的输出；请根据你的硬件设置合适的生成预算。

JustRL II 配方（摘要）

该博客使用带 critic 的 GRPO 配方训练此检查点：

数据：一个三阶段 pipeline 审计约 10 万个开源数学问题（DAPO-Math、DeepScaleR、DeepMath）的可解性，用独立的强模型解答重新检查标签，并移除该检查点已能 8/8 解决的问题 —— 最终训练集保留 32,412 个问题。
算法：保留 GRPO 的组结构和动态采样；一个学习得到的 value 模型（从 policy 初始化，偏差校准到池通过率）通过 GAE 和长度自适应 λ 提供 token 级 advantage；仅尾部超长控制限制过长的 rollout 长度。
Rollout：每个 prompt 在温度 1.0 下采样 8 次，128k token 预算，用 math-verify 的变体进行评分。

完整规格、消融实验和 critic 诊断请参阅博客。

引用

@misc{justrl2,
  title  = {JustRL II: Scaling Small LLMs to 128K Reasoning with a Critic},
  author = {Pan, Haoxuan and Zhou, Chuyue and Li, Xin and others},
  year   = {2026},
  howpublished = {\url{https://panhaoxuan.notion.site/justrl-ii-scaling-small-llms-to-128k-reasoning-with-a-critic}}
}

译自 OpenBMB · HF · 录于二〇二六年九月七日