cleaner olmo3 script

Author: mnoukhov

scripts/train/rlvr/grpo_olmo3.sh

@@ -6,43 +6,10 @@ dataset_mix="hamishivi/rlvr_orz_math_57k_collected 56878"
 # all evals
 # evals="minerva_math::hamish_zs_reasoning,gsm8k::zs_cot_latex,gsm8k::hamish_zs_reasoning,minerva_math_500::hamish_zs_reasoning,zebralogic::hamish_zs_reasoning,aime::hamish_zs_reasoning,agi_eval_english:0shot_cot::hamish_zs_reasoning,gpqa:0shot_cot::hamish_zs_reasoning,ifeval::hamish_zs_reasoning,popqa::hamish_zs_reasoning,mmlu:cot::hamish_zs_reasoning,alpaca_eval_v3::hamish_zs_reasoning,bbh:cot::hamish_zs_reasoning,mbppplus:0-shot-chat::tulu-thinker,codex_humanevalplus:0-shot-chat-v1::tulu-thinker"
 # math evals
-evals="minerva_math::hamish_zs_reasoning,minerva_math_500::hamish_zs_reasoning,aime::hamish_zs_reasoning"
+evals="minerva_math::hamish_zs_reasoning,minerva_math_500::hamish_zs_reasoning,aime:zs_cot_r1::pass_at_32_2024_temp1,aime:zs_cot_r1::pass_at_32_2025_temp1"
 
-# all I've changed with the checkpoints is the config.json, model_type=olmo3 and architectures is OLMo3ForCausalLM 
-# jacob tulu sft
-# model_name_or_path="/weka/oe-adapt-default/michaeln/olmo3/olmo3_reasoning-anneal-tulu3sft-olmo2hparams__8__1751523764/"
-# midtraining no reasoning
-# model_name_or_path="/weka/oe-adapt-default/michaeln/olmo3/anneal-round1-100B-olmo3_7b_no-reasoning-anneal-3c193128_step47684"
-# midtraining with reasoning
-# model_name_or_path="/weka/oe-adapt-default/michaeln/olmo3/anneal-round1-100B-olmo3_7b_with-reasoning-anneal-9d6f76b0_step47684"
-# micro anneals
-# model_name_or_path="/weka/oe-adapt-default/allennlp/deletable_checkpoint/michaeln/olmo3_microanneal-finemath-643cecc4_step4769-hf"
-
-# model_name_or_path="/weka/oe-training-default/ai2-llm/checkpoints/kylel/baseline-olmo2_7b-928646-anneal-100B-dolma2-round1-alldressed-17b22b3a/step47684-hf"
-# gs_model_name="olmo2-alldressed-midtraingin"
-
-# model_name_or_path="/weka/oe-training-default/ai2-llm/checkpoints/OLMo3-midtraining/anneal-round1-100B-olmo3_7b_with-reasoning-anneal-12T-3d39e871/step47684-hf"
-# gs_model_name="olmo3-midtraining-round1"
-
-# model_name_or_path="/weka/oe-training-default/ai2-llm/checkpoints/OLMo3-midtraining/anneal-round2-100B-olmo3_7b_with-reasoning-anneal-12T-53f443c7/step47684-hf"
-#
-# model_name_or_path="/weka/oe-training-default/ai2-llm/checkpoints/OLMo3-midtraining/anneal-round3-webround2-100B-olmo3_7b_with-reasoning-anneal-12T-302b1ae8/step47684-hf"
-# gs_model_name="olmo3-midtraining-round3"
-
-# model_name_or_path="/weka/oe-training-default/ai2-llm/checkpoints/ianm/decon-anneal-round3-webround2-100B-olmo3_7b_with-reasoning-anneal-12T-fc803782/step47684-hf"
-# gs_model_name="olmo3-midtraining-round3-decon"
-
-# model_name_or_path="/weka/oe-training-default/ai2-llm/checkpoints/OLMo3-midtraining/anneal-round4-100B-olmo3_7b-anneal-decon-12T-081e9449/step47684-hf"
-# gs_model_name="olmo3-midtraining-round4"
-
-# model_name_or_path="/weka/oe-adapt-default/jacobm/checkpoints/olmo2-7B-sft/olmo3-hparam-search/olmo3-12t-r3-1e-4-2_epochs-olmo2-tulu3-mix-num_3"
-# gs_model_name="olmo3-midtraining-round3-jacobsft-num3"
-
-# model_name_or_path="/weka/oe-training-default/ai2-llm/checkpoints/OLMo3-midtraining/anneal-round5-100B-olmo3_7b-anneal-decon-12T-00bb6023/step47684-hf"
-# gs_model_name="olmo3-midtraining-round5"
-
-model_name_or_path="/weka/oe-adapt-default/jacobm/checkpoints/olmo2-7B-sft/olmo3-hparam-search/olmo3-12t-r5-100b-olmo2-tulu3-mix-num_3/"
-gs_model_name="olmo3-midtraining-round5-jacobsft-mix3"
+model_name_or_path="/weka/oe-training-default/ai2-llm/checkpoints/OLMo3-midtraining/anneal-round5-100B-olmo3_7b-anneal-decon-12T-00bb6023/step47684-hf"
+gs_model_name="olmo3-midtraining-round5"
 
 exp_name="grpo_mathonly_1m_${gs_model_name}"
 EXP_NAME=${EXP_NAME:-${exp_name}}
@@ -60,7 +27,7 @@ python mason.py \
     --workspace ai2/tulu-thinker \
     --priority high \
     --pure_docker_mode \
-    --image michaeln/open_instruct_dev_uv_olmo3 \
+    --image michaeln/open_instruct_olmo3 \
     --preemptible \
     --num_nodes 2 \
     --env VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 \
@@ -75,10 +42,10 @@ python open_instruct/grpo_fast.py \
     --exp_name ${EXP_NAME} \
     --beta 0.0 \
     --num_samples_per_prompt_rollout 16 \
-    --num_unique_prompts_rollout 64 \
+    --num_unique_prompts_rollout 128 \
     --num_mini_batches 4 \
     --num_epochs 1 \
-    --learning_rate 5e-7 \
+    --learning_rate 1e-6 \
     --per_device_train_batch_size 1 \
     --kl_estimator kl3 \
     --dataset_mixer_list ${dataset_mix} \
@@ -101,9 +68,9 @@ python open_instruct/grpo_fast.py \
     --lr_scheduler_type constant \
     --apply_verifiable_reward true \
     --seed 1 \
-    --num_evals 5 \
-    --save_freq 100 \
-    --checkpoint_state_freq 50 \
+    --local_eval_every 25 \
+    --save_freq 25 \
+    --checkpoint_state_freq 25 \
     --gradient_checkpointing \
     --with_tracking \
     --vllm_enable_prefix_caching \