diff --git a/CHANGELOG.md b/CHANGELOG.md index 8d9fbcc40a..bd92b471eb 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/scripts/train/olmo3/32b_rlzero_math_olmocore.sh b/scripts/train/olmo3/32b_rlzero_math_olmocore.sh new file mode 100755 index 0000000000..3d6346aea7 --- /dev/null +++ b/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 "$@" diff --git a/scripts/train/olmo3/32b_think_rl_olmocore.sh b/scripts/train/olmo3/32b_think_rl_olmocore.sh new file mode 100755 index 0000000000..ead28cb15b --- /dev/null +++ b/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