SaifPunjwani/jrl-checkpoints

JRL and MR-ME checkpoint release

This release ships eight reasoning-model checkpoints together, covering Qwen3-1.7B, Qwen3-4B, and Ministral-3-3B, with GPU and TPU/JAX exports.

Each scorecard must refer to one immutable checkpoint revision and its own evaluation configuration. Results from different checkpoints or evaluation protocols must not be combined into a single scorecard.

Models

Model	Release label	Checkpoint
Qwen3-1.7B	MR-ME	qwen3-1.7b-mrme-ckpt1
Qwen3-1.7B	JRL	qwen3-1.7b-jrl-ckpt2
Qwen3-1.7B	Long-DAPO	qwen3-1.7b-long-dapo-ckpt5
Qwen3-1.7B	DAPO	qwen3-1.7b-dapo
Qwen3-4B	JRL	qwen3-4b-jrl-ckpt3
Qwen3-4B	MR-ME	qwen3-4b-mrme
Ministral-3-3B	JRL	ministral-3-3b-jrl-ckpt4
Ministral-3-3B	MR-ME	ministral-3-3b-mrme

In the method terminology, DAPO is the correctness-only baseline; JRL separates a novelty-rewarded Explorer from the final Main model; MR-ME extends that design to multiple explorers and rounds. The table identifies the release entries. These names alone do not establish a particular export's training recipe, update budget, explorer schedule, or training hardware; those details require its associated run records.

Directory layout

Each entry has gpu/ and tpu/ subfolders. These are loading formats for the corresponding checkpoint, not separate benchmark candidates. The current repository contains inference exports rather than resumable training state. It does not bundle optimizer state, training logs, evaluation outputs, or generated trajectories.

Available exports

Folder	Architecture	GPU export	TPU/JAX export
qwen3-1.7b-mrme-ckpt1	Qwen3-1.7B	safetensors	framework-neutral safetensors
qwen3-1.7b-jrl-ckpt2	Qwen3-1.7B	sharded safetensors	Flax MsgPack
qwen3-4b-jrl-ckpt3	Qwen3-4B	sharded safetensors	Flax MsgPack
qwen3-4b-mrme	Qwen3-4B	safetensors	Flax MsgPack
ministral-3-3b-jrl-ckpt4	Ministral-3-3B	safetensors	Flax MsgPack
ministral-3-3b-mrme	Ministral-3-3B	safetensors	Flax MsgPack
qwen3-1.7b-long-dapo-ckpt5	Qwen3-1.7B	sharded safetensors	Flax MsgPack
qwen3-1.7b-dapo	Qwen3-1.7B	safetensors	framework-neutral safetensors

The gpu/ subfolders are directly loadable with Transformers. The tpu/ subfolders provide JAX-oriented parameter exports. A TPU/JAX serialization specifies how weights are loaded; it does not identify the hardware used to train them.

Transformers inference

Install current inference dependencies:

pip install "transformers>=5.16.1" "mistral-common>=1.11.7" torch

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

repo_id = "SaifPunjwani/jrl-checkpoints"
subfolder = "qwen3-1.7b-mrme-ckpt1/gpu"

device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.bfloat16 if device == "cuda" else torch.float32

tokenizer = AutoTokenizer.from_pretrained(repo_id, subfolder=subfolder)
model = AutoModelForCausalLM.from_pretrained(
    repo_id,
    subfolder=subfolder,
    dtype=dtype,
).to(device)

messages = [{"role": "user", "content": "Compute 2+2. Give only the number."}]
inputs = tokenizer.apply_chat_template(
    messages,
    tokenize=True,
    add_generation_prompt=True,
    return_tensors="pt",
    return_dict=True,
).to(device)

output = model.generate(**inputs, max_new_tokens=64, do_sample=False)
print(tokenizer.decode(
    output[0, inputs["input_ids"].shape[1]:],
    skip_special_tokens=True,
))

Change subfolder to any other gpu/ path in the table.

TPU/JAX loading

For any export containing tpu/flax_model.msgpack, restore the flat Hugging Face-named parameter dictionary from that file:

from flax.serialization import msgpack_restore
from huggingface_hub import hf_hub_download

path = hf_hub_download(
    "SaifPunjwani/jrl-checkpoints",
    "qwen3-1.7b-jrl-ckpt2/tpu/flax_model.msgpack",
)
with open(path, "rb") as handle:
    flat_hf_named_params = msgpack_restore(handle.read())

Checkpoint 1 and the DAPO reference use framework-neutral safetensors in their tpu/ folders; each includes load_params.py for loading them as JAX arrays. The MsgPack exports are flat parameter dictionaries rather than native FlaxAutoModelForCausalLM directory layouts. For the standard Hugging Face generation API, use the corresponding gpu/ subfolder.

End-to-end JAX/TPU inference

The repository now includes a self-contained, inference-only jax_runtime/ loader for all eight TPU exports. It restores either storage format, validates and maps the Hugging Face-named tensors, shards parameters over all visible TPU devices, and performs autoregressive decoding with a KV cache:

python -m pip install --upgrade "jax[tpu]"
python -m pip install -r jax_runtime/requirements.txt
python -m jax_runtime.smoke_generate \
  --folder qwen3-1.7b-jrl-ckpt2 \
  --max-new-tokens 16

Change --folder to any checkpoint listed above. The runtime downloads only that checkpoint's tpu/ directory and applies its own tokenizer and chat template. It is intended for loading verification and small evaluations rather than high-throughput serving. See jax_runtime/README.md for local-path usage.

Evaluation protocol

The earlier release documentation specifies the following math-evaluation reference protocol. Exact reproduction must use the configuration attached to the evaluated checkpoint and generation pool; these defaults do not establish that every release entry was evaluated with identical settings.

• Use the checkpoint's own tokenizer and chat template; enable thinking mode for Qwen3 thinking evaluations.
• System prompt: Please reason step by step, and put your final answer within \\boxed{}.
• Sampling: temperature 0.6, top-p 0.95, top-k 20, min-p 0.
• Reference context cap: 40,960 tokens. The earlier AIME24/AIME25 evaluations allowed up to 38,912 completion tokens; MATH500 and Minerva used up to 32,768. A run using a different budget must be reported separately.
• Retained samples per problem: 64 for AIME24, AIME25, and Minerva; 32 for MATH500.
• Seeds, stop conditions, engine version, and model-specific decoding settings must be taken from the individual evaluation record. In particular, do not assume that Ministral's model-card defaults equal the Qwen3 configuration.

Final answers are extracted from completions and checked against the reference answers with the evaluation's pinned math-answer extractor and verifier. Keep raw completions, correctness labels, finish reasons, sample IDs, checkpoint revision, and evaluation configuration together so the metrics can be recomputed. A short inference example above is a loading demonstration, not a benchmark or a long-form generation-quality audit.

avg@n and benchmark means

For P problems with exactly n retained samples each:

avg@n (%) = 100 * total_correct / (P * n)

With unequal sample counts, average each problem's accuracy rather than silently weighting problems by their number of completions.

The earlier release's four-benchmark mean is the arithmetic mean of AIME24, AIME25, MATH500, and Minerva avg@n. If AIME26 is included, label the result explicitly as a five-benchmark mean and use protocol-compatible evaluations of the same checkpoint. The two aggregate definitions are not interchangeable.

AIME24 pass@k and trajectory evidence

For a problem with c correct samples among n, estimate pass@k as:

1 - C(n - c, k) / C(n, k)

Use zero for C(n - c, k) when n - c < k, and require k <= n. Average this quantity across problems and multiply by 100 for a percentage. For n = 64, pass@64 is the fraction of problems with at least one correct generation. Compute intermediate values at k = 1, 2, 4, 8, 16, 32, 64 from that same sample pool.

The earlier release associated AIME24 pass@k evidence with checkpoints 1, 2, and 5. The current repository tree does not include a pass-at-k-trajectories/ directory; this README does not imply that those archives are bundled here. DAPO-prompt trajectory libraries, when supplied separately, are not substitutes for the AIME24 generations used to calculate AIME24 pass@k.