openbmb / JustRL-II-base-model

JustRL-II Base Model

This is the RL initialization checkpoint used in the blog JustRL II: Scaling Small LLMs to 128K Reasoning with a Critic ( 中文版 ).

It is the starting point of the mathematical-reasoning case study in that blog — the checkpoint that every run there (the standard-GRPO baseline, the full JustRL II recipe, and all ablations) is initialized from and evaluated against. It is not the post-RL model. We release it so that the blog's data pipeline and training recipe can be reproduced from the exact same starting weights.

What this checkpoint is

• A small language model that has already been trained to produce long chain-of-thought (thinking) responses. It emits its reasoning inside <think> ... </think> before the final answer.
• Before any RL, it scores about 61% on AIME 2025 under the blog's evaluation protocol (average over sampled responses, 128k-token generation budget).
• Starting from this checkpoint, the full JustRL II recipe reaches 81% on AIME 2025 in ~300 RL steps, while a standard GRPO baseline on the same data plateaus around 74% (see Figures 1, 2 and 8 of the blog).
• The blog's difficulty-recalibration step (Stage 3 of the data pipeline) is computed by rolling out this checkpoint 8 times per problem, so the released training set's difficulty tiers refer to this model's pass rates.

Files

Two end-of-sequence ids are configured ( eos_token_id = [1, 130073] ); pass both to your generation call or serving engine. config.json ships with max_position_embeddings = 65536 ; the RL runs in the blog use a 128k-token generation budget — refer to the blog for the long-context serving setup used there.

Usage

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

For evaluation or RL rollouts, serve it with vLLM or SGLang as a standard Llama-architecture model, e.g.

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

Intended use

• Reproducing the JustRL II recipe and its ablations from the same initialization.
• Research on long-CoT RL for small models: credit assignment, critic diagnostics, data difficulty calibration.

This checkpoint has not been aligned for general assistant use and has only been evaluated on mathematical reasoning. It may produce very long outputs; set a generation budget appropriate to your hardware.

The JustRL II recipe (summary)

The blog trains this checkpoint with a critic-equipped GRPO recipe:

• Data : a three-stage pipeline audits ~100k open-source math problems (DAPO-Math, DeepScaleR, DeepMath) for solvability, re-checks labels with independent strong-model solutions, and removes problems this checkpoint already solves 8/8 — leaving 32,412 problems in the training set.
• Algorithm : GRPO's group structure and dynamic sampling are retained; a learned value model (initialized from the policy, bias calibrated to the pool pass rate) provides token-level advantages through GAE with a length-adaptive λ; tail-only overlong control limits excessive rollout length.
• Rollouts : 8 samples per prompt at temperature 1.0 with a 128k-token budget, graded by a variant of math-verify.

See the blog for the full specification, ablations, and critic diagnostics.

Citation

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