JaeWooShin/grpo-group-size

GRPO group-size study: checkpoints

Checkpoints from a study of the group in GRPO. The group of G answers per question that the baseline is computed over is separated from the K answers per question that the update is computed on. Each run trains Qwen3-4B-Base with DAPO-style GRPO on the English, deduplicated and decontaminated part of DAPO-Math-17k (14,410 problems) with verl.

• Study 1: 1,024 answers in every update; questions per step = 1024 / K. Compares K at the same G.
• Study 2: 64 questions per step; update size = 64 × K. Compares G at the same update size.

Recipe: token-mean loss; clip 0.2 / 0.28; AdamW lr 1e-6 with 10 warm-up steps; 4 optimizer updates per step; no KL term; temperature 1.0; 8,192-token responses with DAPO's soft overlong penalty. Validation uses AIME24, AIME25, AMC23 and MATH-500.

Layout

One folder per run, named grpo_qwen3-4b-base_dapo_en_g<G>[_k<K>]_b<questions per step>_r8k_s<seed>:

path	content
<run>/merged/global_step_<N>/	bf16 transformers model after N steps (every 50 steps, at the re-evaluated steps, and at the last step)
<run>/rollouts/steps_<a>-<b>.tar.gz	every training rollout: prompt, answer, reward, one <step>.jsonl per step
<run>/results/	metrics.jsonl (verl's per-step metrics), reeval/ (32 samples per AIME/AMC problem), val_generations/, tensorboard/, logs/

The full FSDP checkpoints (weights, optimizer and dataloader state, sharded for 2 GPUs; to continue training with verl) are in the Storage Bucket JaeWooShin/grpo-group-size-checkpoints: the last one of each run (and, while a run trains, its newest one) at <run>/global_step_<N>/, complete once its upload_complete.json is there.

Loading a checkpoint

from transformers import AutoModelForCausalLM, AutoTokenizer

repo = "JaeWooShin/grpo-group-size"
sub = "grpo_qwen3-4b-base_dapo_en_g8_k4_b64_r8k_s0/merged/global_step_200"
model = AutoModelForCausalLM.from_pretrained(repo, subfolder=sub, torch_dtype="bfloat16")
tokenizer = AutoTokenizer.from_pretrained(repo, subfolder=sub)

To serve a checkpoint with vLLM, download the folder first: hf download JaeWooShin/grpo-group-size --include "<run>/merged/global_step_200/*" --local-dir ./ckpt.

The prompt is one ChatML user turn: the problem, then Please reason step by step, and put your final answer within \boxed{}. Generation stops on <|im_end|>.