JRL and MR-ME checkpoint inference exports

This repository contains standalone model exports for inference. It does not include optimizer or scheduler state, RNG state, experiment-tracking artifacts, 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 contain the same model parameters in a JAX-oriented serialization.

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.

Downloads last month: -; Downloads are not tracked for this model. How to track