ftajwar/maxrl_qwen3_4B_base_polaris_rollouts

MaxRL Qwen3-4B-Base training rollouts (POLARIS math prompts)

Every training rollout from an online RL run, with exact token ids, sampling log-probs, and raw rewards — usable as a replay buffer to study off-policy RL for LLM reasoning completely offline.

The run: Qwen3-4B-Base trained with the maxRL advantage estimator (A = (r - mean)/(mean + eps), group mean over 16 rollouts per prompt; maxRL paper) and a pure REINFORCE loss (L = -A * log pi; no importance ratio, no clipping, no KL), on the POLARIS-53K math prompt set with a binary math-verify reward. 1000 training steps ≈ 4.8 epochs, 256 prompts x 16 rollouts per step, sampled by vLLM at temperature 1.0, max response length 4096, trained on 256 A100s (ALCF Polaris).

Files

One Parquet shard per trainer step: data/step-0001.parquet … data/step-1000.parquet. Each shard has exactly 4096 rows = 256 prompts x 16 rollouts, in the exact order the trainer consumed them (verified row-identical to the training dumps). Rows [16k, 16k+16) of a shard are the 16 rollouts of one prompt and share a uid. 4.096M rows total.

Schema

Tokenizer: Qwen3 (vocab 151936), identical across Qwen3 sizes. Note Qwen3-Base uses eos = pad = 151643 (not the chat <|im_end|> 151645). ~1% of rows are length-truncated at 4096 tokens and have no trailing EOS. prompt_token_ids + response_token_ids reconstructs the exact unpadded input_ids; masks are all-ones over the stored lengths (only real tokens are stored). Never re-tokenize the strings — decode/encode is not round-trip-exact; the ids are the ground truth.

Load

from datasets import load_dataset

# everything (streams fine too: streaming=True)
ds = load_dataset("ftajwar/maxrl_qwen3_4B_base_polaris_rollouts", split="train")

# or just a step range
ds = load_dataset(
    "ftajwar/maxrl_qwen3_4B_base_polaris_rollouts",
    data_files=[f"data/step-{s:04d}.parquet" for s in range(1, 51)],
    split="train",
)

Producing run (exact hyperparameters)

Qwen3-4B-Base; AdamW lr 1e-6 (constant), grad-clip 0.3, batch 256 prompts x n=16, 1 PPO epoch (each sample used once, sequential minibatches, no shuffle), gradient checkpointing, bf16 FSDP; vLLM rollouts T=1.0 / top-p 1.0, max prompt 1024, max response 4096; binary math-verify reward; no KL penalty, no entropy bonus, balance_batch=False. The advantage for each row is a pure function of (score, uid-group): A = (r - mean_g) / (mean_g + 1e-6).

Because the loss is plain REINFORCE, a row's gradient contribution depends only on (tokens, advantage, current policy) — this dataset therefore contains everything needed to replay the run offline, step by step, or to resample it into new off-policy schedules (the rollout_log_probs are the behavior-policy log-probs for importance-sampling corrections).

Licenses / provenance

Model outputs from Qwen/Qwen3-4B-Base (Apache-2.0). Prompts from the POLARIS-53K math training set; see POLARIS for upstream prompt provenance.