Add on-policy GSM8K data loop

Author: TimurTaepov

bash/20260529_gsm8k_mc_sampled.sh

@@ -0,0 +1,98 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+# Current implementation:
+#   on-policy loop = live GSM8K candidate generation -> generated parquet -> VERL/GRPO train
+#   -> roll out the saved HF actor checkpoint -> repeat.
+#
+# This is intentionally different from the older static sampled-parquet launcher:
+#   bash/20260327_gsm8k_mc_sampled.sh
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
+BUILDER_DIR="$REPO_ROOT/reliable-gsm8k-builder"
+
+DATE_TAG=${DATE_TAG:-20260529}
+GPU_IDS=${GPU_IDS:-}
+
+BASE_MODEL=${BASE_MODEL:-Qwen/Qwen2.5-3B-Instruct}
+GENERATOR_PROFILE=${GENERATOR_PROFILE:-qwen25_3b}
+INFERENCE_PROFILE=${INFERENCE_PROFILE:-sample_balanced}
+TRAIN_DATASET=${TRAIN_DATASET:-mc_onecorrect}  # mc_onecorrect | mc_allwrong | oe
+
+ITERATIONS=${ITERATIONS:-3}
+NUM_SAMPLES=${NUM_SAMPLES:-1000}
+SPLIT=${SPLIT:-train}
+VAL_FILE=${VAL_FILE:-$HOME/data/gsm8k/test.parquet}
+
+PROJECT_NAME=${PROJECT_NAME:-multiple_choice_question_study}
+IMPLEMENTATION_TAG=${IMPLEMENTATION_TAG:-gsm8k_mc_sampled_onpolicy_livegen_parseronly_grpo}
+RUN_PREFIX=${RUN_PREFIX:-${DATE_TAG}_${IMPLEMENTATION_TAG}_${GENERATOR_PROFILE}_${TRAIN_DATASET}_n${NUM_SAMPLES}_it${ITERATIONS}}
+
+OUTPUT_ROOT=${OUTPUT_ROOT:-$BUILDER_DIR/runs_on_policy}
+CHECKPOINT_ROOT=${CHECKPOINT_ROOT:-$REPO_ROOT/checkpoints/on_policy}
+LOG_DIR=${LOG_DIR:-$REPO_ROOT/logs}
+LOG_FILE=${LOG_FILE:-$LOG_DIR/${RUN_PREFIX}.log}
+
+TOTAL_EPOCHS=${TOTAL_EPOCHS:-1}
+SAVE_FREQ=${SAVE_FREQ:-999999}  # positive; VERL still saves on final step
+TEST_FREQ=${TEST_FREQ:-1}
+TRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-1024}
+PPO_MINI_BATCH_SIZE=${PPO_MINI_BATCH_SIZE:-1024}
+ACTOR_MICRO_BATCH_SIZE=${ACTOR_MICRO_BATCH_SIZE:-8}
+ROLLOUT_N=${ROLLOUT_N:-8}
+LEARNING_RATE=${LEARNING_RATE:-1e-6}
+
+USE_JUDGE=${USE_JUDGE:-0}
+JUDGE_PROFILE=${JUDGE_PROFILE:-judgelm_7b}
+DRY_RUN=${DRY_RUN:-0}
+
+mkdir -p "$LOG_DIR"
+
+CMD=(
+  python "$BUILDER_DIR/run_on_policy_loop.py"
+  --run-prefix "$RUN_PREFIX"
+  --iterations "$ITERATIONS"
+  --num-samples "$NUM_SAMPLES"
+  --split "$SPLIT"
+  --base-model "$BASE_MODEL"
+  --generator-profile "$GENERATOR_PROFILE"
+  --inference-profile "$INFERENCE_PROFILE"
+  --train-dataset "$TRAIN_DATASET"
+  --output-root "$OUTPUT_ROOT"
+  --checkpoint-root "$CHECKPOINT_ROOT"
+  --val-file "$VAL_FILE"
+  --project-name "$PROJECT_NAME"
+  --total-epochs "$TOTAL_EPOCHS"
+  --save-freq "$SAVE_FREQ"
+  --test-freq "$TEST_FREQ"
+  --train-batch-size "$TRAIN_BATCH_SIZE"
+  --ppo-mini-batch-size "$PPO_MINI_BATCH_SIZE"
+  --actor-micro-batch-size "$ACTOR_MICRO_BATCH_SIZE"
+  --rollout-n "$ROLLOUT_N"
+  --learning-rate "$LEARNING_RATE"
+)
+
+if [[ -n "$GPU_IDS" ]]; then
+  export CUDA_VISIBLE_DEVICES="$GPU_IDS"
+  CMD+=(--gpu-ids "$GPU_IDS")
+fi
+
+if [[ "$USE_JUDGE" == "1" ]]; then
+  CMD+=(--use-judge --judge-profile "$JUDGE_PROFILE")
+fi
+
+if [[ "$DRY_RUN" == "1" ]]; then
+  CMD+=(--dry-run)
+fi
+
+echo "[20260529_gsm8k_mc_sampled] implementation=$IMPLEMENTATION_TAG"
+echo "[20260529_gsm8k_mc_sampled] run_prefix=$RUN_PREFIX"
+echo "[20260529_gsm8k_mc_sampled] base_model=$BASE_MODEL"
+echo "[20260529_gsm8k_mc_sampled] generator_profile=$GENERATOR_PROFILE inference_profile=$INFERENCE_PROFILE"
+echo "[20260529_gsm8k_mc_sampled] train_dataset=$TRAIN_DATASET iterations=$ITERATIONS num_samples=$NUM_SAMPLES"
+echo "[20260529_gsm8k_mc_sampled] output_root=$OUTPUT_ROOT"
+echo "[20260529_gsm8k_mc_sampled] checkpoint_root=$CHECKPOINT_ROOT"
+echo "[20260529_gsm8k_mc_sampled] CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-<all>}"
+echo "[20260529_gsm8k_mc_sampled] log_file=$LOG_FILE"
+
+PYTHONUNBUFFERED=1 VLLM_USE_FLASHINFER_SAMPLER=1 "${CMD[@]}" 2>&1 | tee "$LOG_FILE"

reliable-gsm8k-builder/README.md

@@ -93,6 +93,17 @@ python run_build.py \
   --generator-profile qwen25_3b
 ```
 
+Use a local checkpoint/model path as the generator:
+
+```bash
+python run_build.py \
+  --run-id from-checkpoint \
+  --split train \
+  --num-samples 100 \
+  --generator-profile qwen25_3b \
+  --generator-model-path /path/to/actor/huggingface
+```
+
 ## Multi-GPU
 
 If multiple GPUs are visible and either the generator or the optional judge uses a local Transformers model,
@@ -160,6 +171,7 @@ Each run writes to `runs/<run_id>/`:
 - `datasets/gsm8k_oe_<generator>/<split>.jsonl`
 - `datasets/gsm8k_mc_onecorrect_<generator>/<split>.jsonl`
 - `datasets/gsm8k_mc_allwrong_<generator>/<split>.jsonl`
+- matching `.parquet` files for `items` and every generated dataset
 - `manifest.json`
 
 Every item and dataset row includes both:
@@ -198,3 +210,45 @@ FINAL_ANSWER: ...
 - `gsm8k_mc_allwrong` is a real all-wrong dataset:
   - the prompt tells the model to output `NONE` if no option is correct
   - the ground truth is always `NONE`
+
+## On-Policy Loop
+
+`run_on_policy_loop.py` automates:
+
+1. generate `N` fresh data samples from the current model
+2. train with VERL/GRPO on the generated parquet
+3. find the saved HF actor checkpoint
+4. use that checkpoint as the generator for the next iteration
+
+Example dry run:
+
+```bash
+python run_on_policy_loop.py \
+  --run-prefix onpol-smoke \
+  --iterations 2 \
+  --num-samples 100 \
+  --base-model Qwen/Qwen2.5-3B-Instruct \
+  --generator-profile qwen25_3b \
+  --inference-profile sample_balanced \
+  --train-dataset mc_onecorrect \
+  --gpu-ids 0,1,2 \
+  --dry-run
+```
+
+Real run:
+
+```bash
+python run_on_policy_loop.py \
+  --run-prefix onpol-qwen \
+  --iterations 3 \
+  --num-samples 1000 \
+  --base-model Qwen/Qwen2.5-3B-Instruct \
+  --generator-profile qwen25_3b \
+  --inference-profile sample_balanced \
+  --train-dataset mc_onecorrect \
+  --val-file ~/data/gsm8k/test.parquet \
+  --gpu-ids 0,1,2
+```
+
+The loop stores generated data under `runs_on_policy/<run>-iterXX-data/` and checkpoints under
+`../checkpoints/on_policy/<run>-iterXX-train/` by default.

reliable-gsm8k-builder/run_build.py

@@ -53,6 +53,8 @@ def _spawn_multi_gpu_workers(args: argparse.Namespace, *, script_path: Path, gpu
                 str(gpu_id),
                 "--disable-multi-gpu",
             ]
+            if args.generator_model_path is not None:
+                cmd.extend(["--generator-model-path", args.generator_model_path])
             if args.judge_profile is not None:
                 cmd.extend(["--judge-profile", args.judge_profile])
             if args.use_judge:
@@ -86,6 +88,11 @@ def main() -> None:
     parser.add_argument("--num-samples", type=int, default=None, help="Alias for --max-items when you want to run only a small sample.")
     parser.add_argument("--output-root", default=str(SCRIPT_DIR / "runs"), help="Directory that will contain run outputs.")
     parser.add_argument("--generator-profile", default="qwen25_3b", help="Built-in generator profile name.")
+    parser.add_argument(
+        "--generator-model-path",
+        default=None,
+        help="Optional model/checkpoint path overriding the selected generator profile's model/tokenizer.",
+    )
     parser.add_argument("--inference-profile", default="greedy", help="Built-in inference profile name.")
     parser.add_argument("--judge-profile", default="judgelm_7b", help="Built-in judge profile name if --use-judge is enabled.")
     parser.add_argument("--use-judge", action="store_true", help="Use a judge model to verify generated numeric answers. Disabled by default.")
@@ -162,6 +169,7 @@ def main() -> None:
         max_items=args.effective_max_items,
         output_root=output_root,
         generator_profile_name=args.generator_profile,
+        generator_model_path=args.generator_model_path,
         inference_profile_name=args.inference_profile,
         judge_profile_name=args.judge_profile,
         judge_max_tokens=args.judge_max_tokens,