Add OLMo-core 32B GRPO launch scripts

CHANGELOG.md

@@ -4,6 +4,7 @@ All notable changes to this project will be documented in this file.
 
 
 ### Added
+- Add OLMo-core (FSDP) launch scripts `scripts/train/olmo3/32b_think_rl_olmocore.sh` and `scripts/train/olmo3/32b_rlzero_math_olmocore.sh`, which run 32B GRPO via `open_instruct/grpo.py` instead of the DeepSpeed `grpo_fast.py` (https://github.com/allenai/open-instruct/pull/1726).
 
 ### Changed
 - Change the default generation `temperature` to 1.0 and make `SamplingConfig.temperature` a required field so `StreamingConfig.temperature` is the single source of truth (https://github.com/allenai/open-instruct/pull/1725).

scripts/train/olmo3/32b_rlzero_math_olmocore.sh

@@ -0,0 +1,81 @@
+#!/bin/bash
+# 32B OLMo-core (FSDP) version of 7b_rlzero_math.sh.
+# Uses open_instruct/grpo.py (OLMo-core Trainer) instead of grpo_fast.py (DeepSpeed),
+# and scales the model from 7B to 32B. DeepSpeed args are replaced by FSDP args
+# (--fsdp_shard_degree / --fsdp_num_replicas / --activation_memory_budget).
+
+EXP_NAME="olmo3_32b_rlzero_math_olmocore"
+MODEL_NAME_OR_PATH="allenai/Olmo-3-1025-32B"
+DATASETS="allenai/Dolci-RLZero-Math-7B 1.0"
+
+LOCAL_EVALS="allenai/aime_2025_openinstruct 1.0"
+LOCAL_EVAL_SPLITS="train"
+
+EVALS="aime:zs_cot_r1::pass_at_32_2024_rlzero,aime:zs_cot_r1::pass_at_32_2025_rlzero"
+
+BEAKER_USER=$(beaker account whoami --format json | jq -r '.[0].name')
+BEAKER_IMAGE="nathanl/open_instruct_auto"
+
+# Check if the first argument starts with the value of $BEAKER_NAME
+if [[ "$1" == "$BEAKER_USER"* ]]; then
+    BEAKER_IMAGE="$1"
+    shift
+fi
+
+uv run mason.py \
+    --task_name ${EXP_NAME} \
+    --cluster ai2/jupiter \
+    --workspace ai2/olmo-instruct \
+    --priority high \
+    --pure_docker_mode \
+    --image ${BEAKER_IMAGE} \
+    --preemptible \
+    --no_auto_dataset_cache \
+    --num_nodes 9 \
+    --env VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 \
+    --gpus 8 \
+    -- source configs/beaker_configs/ray_node_setup.sh \
+\&\& uv run open_instruct/grpo.py \
+    --exp_name ${EXP_NAME} \
+    --beta 0.0 \
+    --no_resampling_pass_rate 0.875 \
+    --active_sampling \
+    --num_samples_per_prompt_rollout 8 \
+    --num_unique_prompts_rollout 32 \
+    --num_mini_batches 1 \
+    --learning_rate 1e-6 \
+    --per_device_train_batch_size 1 \
+    --kl_estimator 2 \
+    --dataset_mixer_list $DATASETS \
+    --dataset_mixer_list_splits train \
+    --dataset_mixer_eval_list $LOCAL_EVALS \
+    --dataset_mixer_eval_list_splits $LOCAL_EVAL_SPLITS \
+    --max_prompt_token_length 2048 \
+    --response_length 16384 \
+    --pack_length 18432 \
+    --model_name_or_path ${MODEL_NAME_OR_PATH} \
+    --chat_template_name olmo_thinker_rlzero \
+    --non_stop_penalty False \
+    --temperature 1.0 \
+    --total_episodes 768000 \
+    --num_learners_per_node 8 8 8 8 \
+    --fsdp_shard_degree 32 \
+    --fsdp_num_replicas 1 \
+    --activation_memory_budget 0.3 \
+    --vllm_num_engines 10 \
+    --vllm_tensor_parallel_size 4 \
+    --lr_scheduler_type constant \
+    --apply_verifiable_reward true \
+    --seed 1 \
+    --local_eval_every 100 \
+    --save_freq 100 \
+    --checkpoint_state_freq 100 \
+    --gradient_checkpointing \
+    --with_tracking \
+    --vllm_enable_prefix_caching \
+    --mask_truncated_completions False \
+    --oe_eval_max_length 32768 \
+    --try_launch_beaker_eval_jobs_on_weka True \
+    --eval_priority high \
+    --oe_eval_tasks $EVALS \
+    --oe_eval_gpu_multiplier 4 "$@"

scripts/train/olmo3/32b_think_rl_olmocore.sh

@@ -0,0 +1,83 @@
+#!/bin/bash
+# OLMo-core (FSDP) version of 32b_think_rl.sh.
+# Uses open_instruct/grpo.py (OLMo-core Trainer) instead of grpo_fast.py (DeepSpeed).
+# DeepSpeed args (--deepspeed_stage / --deepspeed_zpg / --gather_whole_model) are
+# replaced by FSDP args (--fsdp_shard_degree / --fsdp_num_replicas / --activation_memory_budget).
+
+
+BEAKER_IMAGE=${1:-nathanl/open_instruct_auto}
+export exp_name=test_olmo3_32b_rl_olmocore_run_${RANDOM}
+export data_mix="hamishivi/math_rlvr_mixture_dpo 1.0 hamishivi/code_rlvr_mixture_dpo 1.0 hamishivi/IF_multi_constraints_upto5_filtered_dpo_0625_filter 30186 allenai/rlvr_general_mix-keyword-filtered 21387"
+export model_path=/weka/oe-adapt-default/hamishi/model_checkpoints/olmo3-merge-32b-1e-4-5e-5/olmo3-merge-32b-1e-4-5e-5/
+
+
+uv run python mason.py \
+    --cluster ai2/jupiter \
+    --image $BEAKER_IMAGE \
+    --pure_docker_mode \
+    --workspace ai2/olmo-instruct \
+    --priority urgent \
+    --gs_model_name "sft_olmo3_32b_rl_run_testing" \
+    --preemptible \
+    --num_nodes 28 \
+    --gpus 8 \
+    --max_retries 0 \
+    --no_auto_dataset_cache \
+    --env VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 \
+    --env LD_LIBRARY_PATH=/var/lib/tcpxo/lib64 \
+    --env NCCL_LIB_DIR=/var/lib/tcpxo/lib64 \
+    --env HOSTED_VLLM_API_BASE=http://ceres-cs-aus-447.reviz.ai2.in:8001/v1 \
+    -- source configs/beaker_configs/ray_node_setup.sh \&\& source configs/beaker_configs/code_api_setup.sh \&\& python open_instruct/grpo.py \
+        --exp_name ${exp_name} \
+        --beta 0.0 \
+        --num_samples_per_prompt_rollout 8 \
+        --num_unique_prompts_rollout 64 \
+        --num_mini_batches 1 \
+        --num_epochs 1 \
+        --learning_rate 2e-6 \
+        --per_device_train_batch_size 1 \
+        --output_dir /output \
+        --kl_estimator 2 \
+        --dataset_mixer_list ${data_mix} \
+        --dataset_mixer_list_splits train \
+        --dataset_mixer_eval_list hamishivi/omega-combined 8 allenai/IF_multi_constraints_upto5 8 saurabh5/rlvr_acecoder_filtered 8 hamishivi/tulu_3_rewritten_400k_string_f1_only_v2_nocode_all_filtered_qwen2_5_openthoughts2 4 hamishivi/virtuoussy_multi_subject_rlvr 4 \
+        --dataset_mixer_eval_list_splits train \
+        --max_prompt_token_length 2048 \
+        --response_length 32768 \
+        --pack_length 35840 \
+        --model_name_or_path ${model_path} \
+        --chat_template_name olmo_thinker \
+        --non_stop_penalty False \
+        --mask_truncated_completions False \
+        --temperature 1.0 \
+        --ground_truths_key ground_truth \
+        --sft_messages_key messages \
+        --total_episodes 10000000 \
+        --num_learners_per_node 8 8 8 8 8 8 8 8 8 8 8 8 \
+        --fsdp_shard_degree 32 \
+        --fsdp_num_replicas 3 \
+        --activation_memory_budget 0.1 \
+        --vllm_num_engines 6 \
+        --vllm_tensor_parallel_size 8 \
+        --lr_scheduler_type constant \
+        --apply_verifiable_reward true \
+        --seed 1 \
+        --local_eval_every 50 \
+        --save_freq 25 \
+        --eval_priority urgent \
+        --try_launch_beaker_eval_jobs_on_weka True \
+        --gradient_checkpointing \
+        --with_tracking \
+        --llm_judge_model hosted_vllm/Qwen/Qwen3-32B \
+        --llm_judge_timeout 600 \
+        --llm_judge_max_tokens 2048 \
+        --llm_judge_max_context_length 32768 \
+        --code_api_url https://p9f1719l7f.execute-api.us-west-2.amazonaws.com/prod/test_program \
+        --code_pass_rate_reward_threshold 0.99 \
+        --oe_eval_max_length 32768 \
+        --checkpoint_state_freq 100 \
+        --backend_timeout 1200 \
+        --active_sampling \
+        --oe_eval_beaker_image oe-eval-beaker/oe_eval_olmo2_retrofit_auto \
+        --oe_eval_tasks mmlu:cot::hamish_zs_reasoning_deepseek,bbh:cot::hamish_zs_reasoning_deepseek_v2,gpqa:0shot_cot::qwen3-instruct,zebralogic::hamish_zs_reasoning_deepseek,agi_eval_english:0shot_cot::hamish_zs_reasoning_deepseek,omega_500:0-shot-chat_deepseek,aime:zs_cot_r1::pass_at_32_2024_deepseek,aime:zs_cot_r1::pass_at_32_2025_deepseek,codex_humanevalplus:0-shot-chat::tulu-thinker_deepseek,mbppplus:0-shot-chat::tulu-thinker_deepseek,livecodebench_codegeneration::tulu-thinker_deepseek,alpaca_eval_v3::hamish_zs_reasoning_deepseek,ifeval::hamish_zs_reasoning_deepseek \
+        --vllm_enforce_eager