| --- |
| library_name: transformers |
| pipeline_tag: text-generation |
| language: |
| - en |
| - zh |
| tags: |
| - reasoning |
| - long-cot |
| - math |
| - reinforcement-learning |
| - grpo |
| - justrl |
| --- |
| |
| # 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**](https://panhaoxuan.notion.site/justrl-ii-scaling-small-llms-to-128k-reasoning-with-a-critic) |
| ([中文版](https://panhaoxuan.notion.site/justrl-ii-small-llms-to-128k-reasoning-with-a-critic-cn)). |
|
|
| 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 |
|
|
| | File | Notes | |
| | --- | --- | |
| | `pytorch_model.bin` | bf16 weights, single file | |
| | `config.json` | Llama-style architecture (`LlamaForCausalLM`), loads with stock `transformers` | |
| | `tokenizer.json`, `tokenizer_config.json`, `special_tokens_map.json` | tokenizer | |
| | `chat_template.jinja` | chat template with thinking-mode support (`enable_thinking`) | |
|
|
| 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 |
|
|
| ```python |
| 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. |
|
|
| ```bash |
| 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 |
|
|
| ```bibtex |
| @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}} |
| } |
| ``` |
|
|
