Changes to make DPO run faster

open_instruct/utils.py

@@ -664,13 +664,15 @@ def get_wandb_tags() -> list[str]:
     if "GIT_COMMIT" in os.environ:
         git_commit = os.environ["GIT_COMMIT"]
         tags.append(f"commit: {git_commit}")
-        # try finding the pull request number on github
-        prs = requests.get(f"https://api.github.com/search/issues?q=repo:allenai/open-instruct+is:pr+{git_commit}")
-        if prs.status_code == 200:
-            prs = prs.json()
-            if len(prs["items"]):
-                pr = prs["items"][0]
-                tags.append(f"pr: {pr['number']}")
+        try:
+            prs = requests.get(f"https://api.github.com/search/issues?q=repo:allenai/open-instruct+is:pr+{git_commit}")
+            if prs.status_code == 200:
+                prs = prs.json()
+                if len(prs["items"]):
+                    pr = prs["items"][0]
+                    tags.append(f"pr: {pr['number']}")
+        except requests.exceptions.ConnectionError as e:
+            logger.warning(f"Failed to fetch PR information from GitHub API: {e}")
     if "GIT_BRANCH" in os.environ:
         tags.append(f"branch: {os.environ['GIT_BRANCH']}")
     tags = [tag[:64] for tag in tags if len(tag) > 64]
@@ -1130,8 +1132,8 @@ def launch_ai2_evals_on_weka(
     oe_eval_tasks: list[str] | None = None,
     stop_strings: list[str] | None = None,
     gs_bucket_path: str | None = None,
-    eval_priority: str | None = "normal",
-    eval_workspace: str | None = "ai2/tulu-3-results",
+    eval_priority: str | None = "urgent",
+    eval_workspace: str | None = "ai2/olmo-instruct",
     beaker_image: str | None = None,
     oe_eval_gpu_multiplier: int | None = None,
 ) -> None:

scripts/train/olmo3/32b_dpo_smoke_test.sh

@@ -36,6 +36,8 @@ uv run python mason.py \
    --max_train_samples 150000 \
    --dataset_skip_cache \
    --zero_stage 3 \
+   --zero_hpz_partition_size 1 \
+   --offload_optimizer True \
    --ref_logprobs_cache_dir "/filestore/.cache/" \
    --concatenated_forward False \
    --max_seq_length 16384 \

scripts/train/olmo3/scottg_dpo_sweep_on_unmerged_sft.sh

@@ -0,0 +1,63 @@
+BEAKER_IMAGE=$1
+
+MODEL_NAME=/weka/oe-adapt-default/saumyam/checkpoints/olmo2-7B-sft/rl-sft/olmo3-32b-SFT-5e-5/step10790-hf
+NUM_NODES=16
+# for LR in 8e-8 7e-8 6e-8 9e-8 1e-7 5e-8 2e-7
+for LR in 1e-4
+do
+    EXP_NAME="olmo3-32b-5e10790-DPO-deltas-10k-${LR}-3"
+    uv run python mason.py \
+        --cluster ai2/augusta \
+        --gs_model_name olmo3-32b-SFT-5e-5-step10790 \
+        --workspace ai2/olmo-instruct \
+        --priority urgent \
+        --max_retries 5 \
+	--preemptible \
+        --image scottg/open_instruct_dev_dpo_faster --pure_docker_mode \
+        --env NCCL_LIB_DIR=/var/lib/tcpxo/lib64 \
+        --env LD_LIBRARY_PATH=/var/lib/tcpxo/lib64:$LD_LIBRARY_PATH \
+        --env NCCL_PROTO=Simple,LL128 \
+        --env NCCL_TUNER_CONFIG_PATH=/var/lib/tcpxo/lib64/a3plus_tuner_config_ll128.textproto \
+        --env NCCL_SHIMNET_GUEST_CONFIG_CHECKER_CONFIG_FILE=/var/lib/tcpxo/lib64/a3plus_guest_config_ll128.textproto \
+        --num_nodes $NUM_NODES \
+        --budget ai2/oe-adapt \
+        --no_auto_dataset_cache \
+        --gpus 8 -- source /var/lib/tcpxo/lib64/nccl-env-profile.sh \&\& accelerate launch \
+        --mixed_precision bf16 \
+        --num_processes 8 \
+        --use_deepspeed \
+        --deepspeed_config_file configs/ds_configs/stage3_no_offloading_accelerate.conf \
+        --deepspeed_multinode_launcher standard \
+        open_instruct/dpo_tune_cache.py \
+        --exp_name $EXP_NAME \
+        --model_name_or_path $MODEL_NAME \
+        --tokenizer_name $MODEL_NAME \
+        --use_slow_tokenizer False \
+        --dataset_mixer_list allenai/olmo-3-preference-mix-deltas_reasoning-scottmix-DECON-keyword-filtered 1.0 \
+        --max_train_samples 10000 \
+        --dataset_skip_cache \
+	--zero_stage 3 \
+        --concatenated_forward False \
+        --max_seq_length 16384 \
+        --per_device_train_batch_size 1 \
+        --gradient_accumulation_steps 1 \
+        --learning_rate $LR \
+        --lr_scheduler_type linear \
+        --warmup_ratio 0.1 \
+        --weight_decay 0.0 \
+        --num_train_epochs 1 \
+        --logging_steps 1 \
+        --dpo_loss_type dpo_norm \
+        --dpo_beta 5 \
+        --use_flash_attn \
+        --gradient_checkpointing \
+        --report_to wandb \
+        --chat_template_name olmo123 \
+        --with_tracking \
+        --try_launch_beaker_eval_jobs False \
+        --log_grad_norm True \
+	--ref_logprobs_cache_dir "/filestore/.cache/"
+done
+        # --oe_eval_max_length 32768 \
+        # --oe_eval_tasks "gpqa:0shot_cot::qwen3-instruct,codex_humanevalplus:0-shot-chat::tulu-thinker_deepseek,mbppplus:0-shot-chat::tulu-thinker_deepseek,alpaca_eval_v3::hamish_zs_reasoning_deepseek,ifeval::hamish_zs_reasoning_deepseek,agi_eval_english:0shot_cot::hamish_zs_reasoning_deepseek,omega_500:0-shot-chat_deepseek,minerva_math_500::hamish_zs_reasoning_deepseek,livecodebench_codegeneration::tulu-thinker_deepseek_no_think_tags_lite,aime:zs_cot_r1::pass_at_32_2024_deepseek,aime:zs_cot_r1::pass_at_32_2025_deepseek,zebralogic::hamish_zs_reasoning_deepseek,bbh:cot::hamish_zs_reasoning_deepseek_v2,mmlu:cot::hamish_zs_reasoning_deepseek,popqa::hamish_zs_reasoning_deepseek"
+

scripts/train/olmo3/scottg_dpo_sweep_on_unmerged_sft_8breject.sh

@@ -0,0 +1,64 @@
+BEAKER_IMAGE=$1
+
+MODEL_NAME=/weka/oe-adapt-default/saumyam/checkpoints/olmo2-7B-sft/rl-sft/olmo3-32b-SFT-5e-5/step10790-hf
+NUM_NODES=16
+# for LR in 8e-8 7e-8 6e-8 9e-8 1e-7 5e-8 2e-7
+for LR in 8e-7
+do
+    EXP_NAME="olmo3-32b-5e10790-DPO-deltas-10k-${LR}-8b"
+    uv run python mason.py \
+        --cluster ai2/augusta \
+        --gs_model_name olmo3-32b-SFT-5e-5-step10790 \
+        --workspace ai2/olmo-instruct \
+        --priority urgent \
+        --max_retries 5 \
+	--preemptible \
+        --image scottg/open_instruct_dev_dpo_faster --pure_docker_mode \
+        --env NCCL_LIB_DIR=/var/lib/tcpxo/lib64 \
+        --env LD_LIBRARY_PATH=/var/lib/tcpxo/lib64:$LD_LIBRARY_PATH \
+        --env NCCL_PROTO=Simple,LL128 \
+        --env NCCL_TUNER_CONFIG_PATH=/var/lib/tcpxo/lib64/a3plus_tuner_config_ll128.textproto \
+        --env NCCL_SHIMNET_GUEST_CONFIG_CHECKER_CONFIG_FILE=/var/lib/tcpxo/lib64/a3plus_guest_config_ll128.textproto \
+        --num_nodes $NUM_NODES \
+        --budget ai2/oe-adapt \
+        --no_auto_dataset_cache \
+        --gpus 8 -- source /var/lib/tcpxo/lib64/nccl-env-profile.sh \&\& accelerate launch \
+        --mixed_precision bf16 \
+        --num_processes 8 \
+        --use_deepspeed \
+        --deepspeed_config_file configs/ds_configs/stage3_no_offloading_accelerate.conf \
+        --deepspeed_multinode_launcher standard \
+        open_instruct/dpo_tune_cache.py \
+        --exp_name $EXP_NAME \
+        --model_name_or_path $MODEL_NAME \
+        --tokenizer_name $MODEL_NAME \
+        --use_slow_tokenizer False \
+        --dataset_mixer_list allenai/olmo-3-preference-mix-deltas_reasoning-8b_reject-scottmix-DECON-keyword-filtered 1.0 \
+        --max_train_samples 10000 \
+        --dataset_skip_cache \
+	--zero_stage 3 \
+        --concatenated_forward False \
+        --max_seq_length 16384 \
+        --per_device_train_batch_size 1 \
+        --gradient_accumulation_steps 1 \
+        --learning_rate $LR \
+        --lr_scheduler_type linear \
+        --warmup_ratio 0.1 \
+        --weight_decay 0.0 \
+        --num_train_epochs 1 \
+        --logging_steps 1 \
+        --dpo_loss_type dpo_norm \
+        --dpo_beta 5 \
+        --use_flash_attn \
+        --gradient_checkpointing \
+        --report_to wandb \
+        --chat_template_name olmo123 \
+        --with_tracking \
+        --try_launch_beaker_eval_jobs False \
+        --log_grad_norm True \
+	--ref_logprobs_cache_dir "/filestore/.cache/"
+done
+        # --oe_eval_max_length 32768 \
+        # --oe_eval_tasks "gpqa:0shot_cot::qwen3-instruct,codex_humanevalplus:0-shot-chat::tulu-thinker_deepseek,mbppplus:0-shot-chat::tulu-thinker_deepseek,alpaca_eval_v3::hamish_zs_reasoning_deepseek,ifeval::hamish_zs_reasoning_deepseek,agi_eval_english:0shot_cot::hamish_zs_reasoning_deepseek,omega_500:0-shot-chat_deepseek,minerva_math_500::hamish_zs_reasoning_deepseek,livecodebench_codegeneration::tulu-thinker_deepseek_no_think_tags_lite,aime:zs_cot_r1::pass_at_32_2024_deepseek,aime:zs_cot_r1::pass_at_32_2025_deepseek,zebralogic::hamish_zs_reasoning_deepseek,bbh:cot::hamish_zs_reasoning_deepseek_v2,mmlu:cot::hamish_zs_reasoning_deepseek,popqa::hamish_zs_reasoning_deepseek"
+
+	# --ref_logprobs_cache_dir "/filestore/.cache/" \

scripts/train/olmo3/scottg_dpo_sweep_on_unmerged_sft_jupiter.sh

@@ -0,0 +1,57 @@
+BEAKER_IMAGE=$1
+
+MODEL_NAME=/weka/oe-adapt-default/saumyam/checkpoints/olmo2-7B-sft/rl-sft/olmo3-32b-SFT-5e-5/step10790-hf
+NUM_NODES=16
+# for LR in 8e-8 7e-8 6e-8 9e-8 1e-7 5e-8 2e-7
+for LR in 8e-7
+do
+    EXP_NAME="olmo3-32b-5e10790-DPO-deltas-10k-${LR}-j"
+    uv run python mason.py \
+        --cluster ai2/jupiter \
+        --gs_model_name olmo3-32b-SFT-5e-5-step10790 \
+        --workspace ai2/olmo-instruct \
+        --priority urgent \
+        --max_retries 5 \
+	--preemptible \
+        --image scottg/open_instruct_dev_dpo_faster --pure_docker_mode \
+        --num_nodes $NUM_NODES \
+        --budget ai2/oe-adapt \
+        --no_auto_dataset_cache \
+        --gpus 8 -- accelerate launch \
+        --mixed_precision bf16 \
+        --num_processes 8 \
+        --use_deepspeed \
+        --deepspeed_config_file configs/ds_configs/stage3_no_offloading_accelerate.conf \
+        --deepspeed_multinode_launcher standard \
+        open_instruct/dpo_tune_cache.py \
+        --exp_name $EXP_NAME \
+        --model_name_or_path $MODEL_NAME \
+        --tokenizer_name $MODEL_NAME \
+        --use_slow_tokenizer False \
+        --dataset_mixer_list allenai/olmo-3-preference-mix-deltas_reasoning-scottmix-DECON-keyword-filtered 1.0 \
+        --max_train_samples 10000 \
+        --dataset_skip_cache \
+	--zero_stage 3 \
+        --concatenated_forward False \
+        --max_seq_length 16384 \
+        --per_device_train_batch_size 1 \
+        --gradient_accumulation_steps 1 \
+        --learning_rate $LR \
+        --lr_scheduler_type linear \
+        --warmup_ratio 0.1 \
+        --weight_decay 0.0 \
+        --num_train_epochs 1 \
+        --logging_steps 1 \
+        --dpo_loss_type dpo_norm \
+        --dpo_beta 5 \
+        --use_flash_attn \
+        --gradient_checkpointing \
+        --report_to wandb \
+        --chat_template_name olmo123 \
+        --with_tracking \
+        --try_launch_beaker_eval_jobs False \
+        --log_grad_norm True
+done
+        # --oe_eval_max_length 32768 \
+        # --oe_eval_tasks "gpqa:0shot_cot::qwen3-instruct,codex_humanevalplus:0-shot-chat::tulu-thinker_deepseek,mbppplus:0-shot-chat::tulu-thinker_deepseek,alpaca_eval_v3::hamish_zs_reasoning_deepseek,ifeval::hamish_zs_reasoning_deepseek,agi_eval_english:0shot_cot::hamish_zs_reasoning_deepseek,omega_500:0-shot-chat_deepseek,minerva_math_500::hamish_zs_reasoning_deepseek,livecodebench_codegeneration::tulu-thinker_deepseek_no_think_tags_lite,aime:zs_cot_r1::pass_at_32_2024_deepseek,aime:zs_cot_r1::pass_at_32_2025_deepseek,zebralogic::hamish_zs_reasoning_deepseek,bbh:cot::hamish_zs_reasoning_deepseek_v2,mmlu:cot::hamish_zs_reasoning_deepseek,popqa::hamish_zs_reasoning_deepseek"
+

scripts/train/olmo3/scottg_dpo_sweep_on_unmerged_sft_oldenv.sh

@@ -0,0 +1,65 @@
+BEAKER_IMAGE=$1
+
+MODEL_NAME=/weka/oe-adapt-default/saumyam/checkpoints/olmo2-7B-sft/rl-sft/olmo3-32b-SFT-5e-5/step10790-hf
+NUM_NODES=16
+# for LR in 8e-8 7e-8 6e-8 9e-8 1e-7 5e-8 2e-7
+for LR in 8e-7
+do
+    EXP_NAME="olmo3-32b-5e10790-DPO-deltas-10k-${LR}env2"
+    uv run python mason.py \
+        --cluster ai2/augusta \
+        --gs_model_name olmo3-32b-SFT-5e-5-step10790 \
+        --workspace ai2/olmo-instruct \
+        --priority urgent \
+        --max_retries 5 \
+	--preemptible \
+        --image scottg/open_instruct_dev_dpo_faster --pure_docker_mode \
+        --num_nodes $NUM_NODES \
+        --budget ai2/oe-adapt \
+        --no_auto_dataset_cache \
+        --gpus 8 -- accelerate launch \
+        --mixed_precision bf16 \
+        --num_processes 8 \
+        --use_deepspeed \
+        --deepspeed_config_file configs/ds_configs/stage3_no_offloading_accelerate.conf \
+        --deepspeed_multinode_launcher standard \
+        open_instruct/dpo_tune_cache.py \
+        --exp_name $EXP_NAME \
+        --model_name_or_path $MODEL_NAME \
+        --tokenizer_name $MODEL_NAME \
+        --use_slow_tokenizer False \
+        --dataset_mixer_list allenai/olmo-3-preference-mix-deltas_reasoning-scottmix-DECON-keyword-filtered 1.0 \
+        --max_train_samples 10000 \
+        --dataset_skip_cache \
+	--zero_stage 3 \
+        --concatenated_forward False \
+        --max_seq_length 16384 \
+        --per_device_train_batch_size 1 \
+        --gradient_accumulation_steps 1 \
+        --learning_rate $LR \
+        --lr_scheduler_type linear \
+        --warmup_ratio 0.1 \
+        --weight_decay 0.0 \
+        --num_train_epochs 1 \
+        --logging_steps 1 \
+        --dpo_loss_type dpo_norm \
+        --dpo_beta 5 \
+        --use_flash_attn \
+        --gradient_checkpointing \
+        --report_to wandb \
+        --chat_template_name olmo123 \
+        --with_tracking \
+        --try_launch_beaker_eval_jobs False \
+        --log_grad_norm True \
+	--ref_logprobs_cache_dir "/filestore/.cache/"
+done
+        # --oe_eval_max_length 32768 \
+        # --oe_eval_tasks "gpqa:0shot_cot::qwen3-instruct,codex_humanevalplus:0-shot-chat::tulu-thinker_deepseek,mbppplus:0-shot-chat::tulu-thinker_deepseek,alpaca_eval_v3::hamish_zs_reasoning_deepseek,ifeval::hamish_zs_reasoning_deepseek,agi_eval_english:0shot_cot::hamish_zs_reasoning_deepseek,omega_500:0-shot-chat_deepseek,minerva_math_500::hamish_zs_reasoning_deepseek,livecodebench_codegeneration::tulu-thinker_deepseek_no_think_tags_lite,aime:zs_cot_r1::pass_at_32_2024_deepseek,aime:zs_cot_r1::pass_at_32_2025_deepseek,zebralogic::hamish_zs_reasoning_deepseek,bbh:cot::hamish_zs_reasoning_deepseek_v2,mmlu:cot::hamish_zs_reasoning_deepseek,popqa::hamish_zs_reasoning_deepseek"
+
+
+
+        # --env NCCL_LIB_DIR=/var/lib/tcpxo/lib64 \
+        # --env LD_LIBRARY_PATH=/var/lib/tcpxo/lib64:$LD_LIBRARY_PATH \
+        # --env NCCL_PROTO=Simple,LL128 \
+        # --env NCCL_TUNER_CONFIG_PATH=/var/lib/tcpxo/lib64/a3plus_tuner_config_ll128.textproto \
+        # --env NCCL_SHIMNET_GUEST_CONFIG_CHECKER_CONFIG_FILE=/var/lib/tcpxo/lib64/a3plus_guest_config_ll128.textproto \

scripts/train/olmo3/scottg_dpo_sweep_on_unmerged_sft_overfit.sh

@@ -0,0 +1,62 @@
+BEAKER_IMAGE=$1
+
+MODEL_NAME=/weka/oe-adapt-default/saumyam/checkpoints/olmo2-7B-sft/rl-sft/olmo3-32b-SFT-5e-5/step10790-hf
+NUM_NODES=16
+# for LR in 8e-8 7e-8 6e-8 9e-8 1e-7 5e-8 2e-7
+for LR in 1e-6
+do
+    EXP_NAME="olmo3-32b-5e10790-DPO-deltas-${LR}-overfit"
+    uv run python mason.py \
+        --cluster ai2/augusta \
+        --gs_model_name olmo3-32b-SFT-5e-5-step10790 \
+        --workspace ai2/olmo-instruct \
+        --priority urgent \
+        --max_retries 5 \
+	--preemptible \
+        --image scottg/open_instruct_dev_dpo_faster --pure_docker_mode \
+        --env NCCL_LIB_DIR=/var/lib/tcpxo/lib64 \
+        --env LD_LIBRARY_PATH=/var/lib/tcpxo/lib64:$LD_LIBRARY_PATH \
+        --env NCCL_PROTO=Simple,LL128 \
+        --env NCCL_TUNER_CONFIG_PATH=/var/lib/tcpxo/lib64/a3plus_tuner_config_ll128.textproto \
+        --env NCCL_SHIMNET_GUEST_CONFIG_CHECKER_CONFIG_FILE=/var/lib/tcpxo/lib64/a3plus_guest_config_ll128.textproto \
+        --num_nodes $NUM_NODES \
+        --budget ai2/oe-adapt \
+        --no_auto_dataset_cache \
+        --gpus 8 -- source /var/lib/tcpxo/lib64/nccl-env-profile.sh \&\& accelerate launch \
+        --mixed_precision bf16 \
+        --num_processes 8 \
+        --use_deepspeed \
+        --deepspeed_config_file configs/ds_configs/stage3_no_offloading_accelerate.conf \
+        --deepspeed_multinode_launcher standard \
+        open_instruct/dpo_tune_cache.py \
+        --exp_name $EXP_NAME \
+        --model_name_or_path $MODEL_NAME \
+        --tokenizer_name $MODEL_NAME \
+        --use_slow_tokenizer False \
+        --dataset_mixer_list allenai/olmo-3-preference-mix-deltas_reasoning-scottmix-DECON-keyword-filtered 128 \
+        --dataset_skip_cache \
+	--zero_stage 3 \
+        --concatenated_forward False \
+        --max_seq_length 16384 \
+        --per_device_train_batch_size 1 \
+        --gradient_accumulation_steps 1 \
+        --learning_rate $LR \
+        --lr_scheduler_type linear \
+        --warmup_ratio 0.1 \
+        --weight_decay 0.0 \
+        --num_train_epochs 10 \
+        --logging_steps 1 \
+        --dpo_loss_type dpo_norm \
+        --dpo_beta 5 \
+        --use_flash_attn \
+        --gradient_checkpointing \
+        --report_to wandb \
+        --chat_template_name olmo123 \
+        --with_tracking \
+        --try_launch_beaker_eval_jobs False \
+        --log_grad_norm True
+done
+        # --oe_eval_max_length 32768 \
+        # --oe_eval_tasks "gpqa:0shot_cot::qwen3-instruct,codex_humanevalplus:0-shot-chat::tulu-thinker_deepseek,mbppplus:0-shot-chat::tulu-thinker_deepseek,alpaca_eval_v3::hamish_zs_reasoning_deepseek,ifeval::hamish_zs_reasoning_deepseek,agi_eval_english:0shot_cot::hamish_zs_reasoning_deepseek,omega_500:0-shot-chat_deepseek,minerva_math_500::hamish_zs_reasoning_deepseek,livecodebench_codegeneration::tulu-thinker_deepseek_no_think_tags_lite,aime:zs_cot_r1::pass_at_32_2024_deepseek,aime:zs_cot_r1::pass_at_32_2025_deepseek,zebralogic::hamish_zs_reasoning_deepseek,bbh:cot::hamish_zs_reasoning_deepseek_v2,mmlu:cot::hamish_zs_reasoning_deepseek,popqa::hamish_zs_reasoning_deepseek"
+
+	# --ref_logprobs_cache_dir "/filestore/.cache/" \