
# Best so far, using postfinetune v1 checkpoint, 7/1 gpu split, 28 gen 28 accum, beta 0.01, 2e-6 learning rate linear decay, regular rewards
# Apricot Dream
# https://beaker.allen.ai/orgs/ai2/workspaces/olmocr/work/01KAW51JANNX0FNP0BY0ATJW5E?taskId=01KAW51JAQCRW93GC5S8FW2ZEW&jobId=01KAW51JKR75PP3XT545J9HDQE
# s3://ai2-oe-data/jakep/olmocr-grpo-checkpoints/qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune-multigpu-01KAW51JANNX0FNP0BY0ATJW5E
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v1-266 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--reward_bench 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune \
--seed 1 --importance_sampling_level sequence \
--gradient_accumulation_steps 28 \
--learning_rate 2e-6 \
--preemptible

# Same thing, but using more gradient accumulation steps
# s3://ai2-oe-data/jakep/olmocr-grpo-checkpoints/qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_gradaccum84-multigpu-01KAYH934F58SRRX5VVR4CP04S
# Outputs collapse After step 250
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v1-266 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 2e-6 \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_gradaccum84 \
--preemptible

# Same thing with constant lr schedule
# s3://ai2-oe-data/jakep/olmocr-grpo-checkpoints/qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_constnt-multigpu-01KAYPAR4Z7D142XMWK088B1P4
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v1-266 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--num_generations 28 --gradient_accumulation_steps 28 \
--learning_rate 2e-6 --lr_schedule constant \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_constnt \
--preemptible

# Same thing but with top p = 0.9
# Top p =0.9 definitely has a higher fraction of samples with a zero advantage
# And looks like in loop evals are reliably a few points higher
# But absent tests (headers/footers are likely to be worse)
# s3://ai2-oe-data/jakep/olmocr-grpo-checkpoints/qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_topp0_9-multigpu-01KAYHAB11N6MX1M0R1DDJ353Z
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v1-266 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--top_p 0.9 \
--num_generations 28 --gradient_accumulation_steps 28 \
--learning_rate 2e-6 \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_topp0_9 \
--preemptible

# Same thing with macroavg reward
# s3://ai2-oe-data/jakep/olmocr-grpo-checkpoints/qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_macroavg-multigpu-01KAYPFQ11BYYM4MPJN3VWBBYZ
# Table performance tracks an almost exact same decline as our baseline
# Math performance seems to go up slowler but more steadily
# Other test scores seem to track roughly the same also
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v1-266 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--num_generations 28 --gradient_accumulation_steps 28 \
--learning_rate 2e-6 \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_macroavg \
--preemptible

# Same thing with more data 
# Overall rewards are much lower, but maybe this is because the newly generated data is just harder
# But the model order tests for example pass rate is much lower overall
# Maybe the model just diverges too much after a long amount of training
# s3://ai2-oe-data/jakep/olmocr-grpo-checkpoints/qwen25vl_olmocrv5_arxiv10k_general10k_bigtables_postfinetune-multigpu-01KAYHAN116BZS4RE01Q8N0M0Y
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v1-266 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|general_10k|big_tables" \
--reward_bench 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--num_generations 28 --gradient_accumulation_steps 28 \
--learning_rate 2e-6 \
--name qwen25vl_olmocrv5_arxiv10k_general10k_bigtables_postfinetune \
--preemptible

# YOLO run
# Didn't let this run very long
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v1-266 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo \
--preemptible

# YOLO Run with v2 midtrained
# s3://ai2-oe-data/jakep/olmocr-grpo-checkpoints/qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2-multigpu-01KB45ZE1GQCJAREV7EB4R6D9R
# Saw a steady rise in macroavg reward over a decently long training run
# Pretty high fraction on zero reward cases, as expected from top p set higher
# Math had a slow and steady nice improvement
# Olmocr mix had good performance too
# But table and absense (header and footer) scores do still show a steady down turn over time
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2 \
--preemptible

# Run, trying with table tests broken out into individual bite size tests instead of each cell having multiple conditions it must satisfy
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_tablesplit" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2_tablesplit \
--preemptible

# Best YOLO run but with general dataset instead of arxiv mix
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|general_10k" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--name qwen25vl_olmocrv5_arxiv10k_general10k_postfinetune_yolo_v2 \
--preemptible

# Set of runs similar to YOLO but with TRL nightly various combinations
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--vllm_importance_sampling_correction True --vllm_importance_sampling_mode sequence_mask --vllm_importance_sampling_cap 3.0 \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2_trl_nightly_sequence_mask \
--preemptible

./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level token \
--vllm_importance_sampling_correction True --vllm_importance_sampling_mode token_mask --vllm_importance_sampling_cap 3.0 \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2_trl_nightly_token_mask \
--preemptible

./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level token \
--vllm_importance_sampling_correction True --vllm_importance_sampling_mode token_truncate --vllm_importance_sampling_cap 3.0 \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2_trl_nightly_token_truncate \
--preemptible

./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--vllm_importance_sampling_correction True --vllm_importance_sampling_mode sequence_truncate --vllm_importance_sampling_cap 3.0 \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2_trl_nightly_sequence_truncate \
--preemptible

# TRL nightly with general_10k mix, for a cleaner story
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|general_10k$" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--vllm_importance_sampling_correction True --vllm_importance_sampling_mode sequence_mask --vllm_importance_sampling_cap 3.0 \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--name qwen25vl_olmocrv5_arxiv10k_general10k_postfinetune_yolo_v2_trl_nightly_sequence_mask \
--preemptible

# Runs that are split up by test type to see if we can learn something that way
./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--vllm_importance_sampling_correction True --vllm_importance_sampling_mode sequence_truncate --vllm_importance_sampling_cap 3.0 \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--bench_type_filter table \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2_trl_nightly_sequence_truncate_onlytables \
--preemptible

./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--vllm_importance_sampling_correction True --vllm_importance_sampling_mode sequence_truncate --vllm_importance_sampling_cap 3.0 \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--bench_type_filter math \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2_trl_nightly_sequence_truncate_onlymath \
--preemptible

./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--vllm_importance_sampling_correction True --vllm_importance_sampling_mode sequence_truncate --vllm_importance_sampling_cap 3.0 \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--bench_type_filter absent \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2_trl_nightly_sequence_truncate_onlyabsent \
--preemptible

./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--vllm_importance_sampling_correction True --vllm_importance_sampling_mode sequence_truncate --vllm_importance_sampling_cap 3.0 \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--bench_type_filter format \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2_trl_nightly_sequence_truncate_onlyformat \
--preemptible

./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--vllm_importance_sampling_correction True --vllm_importance_sampling_mode sequence_truncate --vllm_importance_sampling_cap 3.0 \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--bench_type_filter order \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2_trl_nightly_sequence_truncate_onlyorder \
--preemptible

./scripts/train/grpotrainer-beaker-multi-gpu.sh --num-train-gpus 7 --num-generate-gpus 1 --model_name s3://ai2-oe-data/jakep/olmocr/qwen25_vl_olmocrv5_synth10postfinetune_v2-960 \
--train_bench_data_folder /data/jakep/grpo_data_mixes/synth_mix_1225/bench_data \
--jsonl_filter "arxiv|olmocr_mix_documents" \
--reward_bench_macroavg 1.0 --reward_front_matter 1.0 --reward_eos 1.0 \
--beta 0.01 \
--seed 1 --importance_sampling_level sequence \
--vllm_importance_sampling_correction True --vllm_importance_sampling_mode sequence_truncate --vllm_importance_sampling_cap 3.0 \
--num_generations 28 --gradient_accumulation_steps 84 \
--learning_rate 1e-6 --lr_schedule constant \
--top_p 0.9 \
--bench_type_filter footnote \
--name qwen25vl_olmocrv5_arxiv10k_mix10k_postfinetune_yolo_v2_trl_nightly_sequence_truncate_onlyfootnote \
--preemptible