Post-training loss curves and training time estimates for OLMo-2 1B

[BrownianNotion]

Hi team,
Would you be able to share the post-training loss curves and training time estimates for each post-training stage of the OLMo-2 1B model on an 8xH100 cluster? I cannot seem to find these in the OLMo-2 paper.

I've reviewed #485 and the OLMo-3 paper (which mentions a 9-day estimate), but both reference either larger models or more GPUs, so I'm hoping to get figures specific to this setup.

Thanks so much for any guidance!

[hamishivi]

Hi,

SFT: 9 hours on one 1 H100 node (8 gpus).
DPO: 2 hours on one 1 H100 node (8 gpus).
RL: 2000000 episodes took roughly 43 hours on 16 H100s (2 nodes). We did two stages of RL training with same cost (so ~86 hours on 16 gpus all up).

Sharing the curves is a bit harder... I am not sure on the precise way to move out our wandb curves. Ping me again in the new year and I can see what I can do!

[BrownianNotion]

Thanks @hamishivi, will reach out in the new year about Wandb! DPO estimate lines up with that I've seen (perhaps closer to 1.5 hours with deepspeed stage 0). Haven't done any full SFT runs but here are some of my estimates on one 8xH100 node (based on tqdm and running ~200 steps):

• default: (docs/olmo2.md) 8+ hours (lines up with your estimate)
• per_device_batch_size=4, gradient_accumulation_steps=4: (same effective batch size) 5.5-6 hours
• deepspeed stage 0: 3.5-4 hours (per_device_batch_size=2)

It seems using deepspeed stage 0 instead of stage 2 gives huge time savings given that the model already fits in memory.

[BrownianNotion]

Hey @hamishivi, did you have any luck on sharing the wandb curves?

[hamishivi]

SFT: https://wandb.ai/ai2-llm/open_instruct_public/runs/mihc7mab
DPO: https://wandb.ai/ai2-llm/open_instruct_public/runs/k34ppe9q
The RL curves didnt transfer for some reason... working on it.

[BrownianNotion]

Hey @hamishivi , thanks so much!
Any chance you know why the loss magnitudes are so different? Green is olmo, blue is my repro. Was the loss function, reduction method etc. changed? It's still a sum loss on finetune.py as far as I can see though.

[hamishivi]

Hi @BrownianNotion , this would be because we adjusted the metrics logging in the SFT script to normalize by the number of tokens, which makes things more comparable when changing the number of gpus used to train: https://github.com/allenai/open-instruct/blob/main/open_instruct/finetune.py#L860-L861

[hamishivi]

Unfortunately, I wasn't able to copy over the wandb curves for the RL jobs, it just didn't work. Here is some curves from the first RL job:

And for the second:

Note for this second run we only used MATH data so the reward shown is just the MATH reward.

Let me know if seeing a particular curve would be useful, sorry for the inconvenience!

[BrownianNotion]

@hamishivi thank you so much, these are still super helpful! If it's not too much trouble, I'd also love to see the loss curves (total, policy kl) and advantages (min, max, hist, val). Once again, really appreciate the help!

[hamishivi]

Here's the KL and loss curves! Unfortunately we did these runs before we had advantage logging.
Generally, I would expect loss to be flat and kl to slowly increase over time.

Run 1:

Run 2:

[BrownianNotion]

Hey @hamishivi, I ran a few steps of the RLVR stage 1 training, but am getting very different loss magnitudes to your graphs. Did you also change the loss scaling/normalisation from the older script grpo_vllm_thread_ray_gtrl.py to the newer grpo_fast.py? This was run on an older checkout of the repo 5a2407345dea38937d2d5753cadff21c5e8777c8, before the Olmo core refactor of grpo_fast.py.

The command I ran was

python open_instruct/grpo_fast.py \
    --add_bos \
    --apply_verifiable_reward true \
    --beta 0.01 \
    --chat_template_name tulu \
    --dataset_cache_mode hf \
    --dataset_config_hash 7691399bb0 \
    --dataset_config_eval_hash 24054476c2 \
    --dataset_mixer_eval_list allenai/RLVR-GSM-MATH-IF-Mixed-Constraints 16 \
    --dataset_mixer_eval_list_splits train \
    --dataset_mixer_list allenai/RLVR-GSM-MATH-IF-Mixed-Constraints 1.0 \
    --dataset_mixer_list_splits train \
    --deepspeed_stage 2 \
    --deepspeed_zpg 12 \
    --exp_name 0423_grpo_seed_1_lr_5e-7_test_single_node \
    --fused_optimizer True \
    --kl_estimator 3 \
    --learning_rate 5e-7 \
    --local_eval_every 5 \
    --lr_scheduler_type constant \
    --max_prompt_token_length 2048 \
    --model_name_or_path allenai/OLMo-2-0425-1B-DPO \
    --model_revision main \
    --non_stop_penalty \
    --non_stop_penalty_value 0.0 \
    --num_epochs 1 \
    --num_learners_per_node 4 8 \
    --num_mini_batches 2 \
    --num_samples_per_prompt_rollout 16 \
    --num_unique_prompts_rollout 48 \
    --pack_length 4096 \
    --per_device_train_batch_size 1 \
    --remap_verifier ifeval=ifeval_old \
    --response_length 2048 \
    --save_freq 200 \
    --save_traces True \
    --seed 1 \
    --stop_strings "</answer>" \
    --temperature 1.0 \
    --tokenizer_name_or_path allenai/OLMo-2-0425-1B-DPO \
    --tokenizer_revision main \
    --total_episodes 2000000 \
    --verbose False \
    --vllm_gpu_memory_utilization 0.9 \
    --vllm_num_engines 4 \
    --vllm_sync_backend nccl \
    --vllm_tensor_parallel_size 1 \
    --with_tracking

[hamishivi]

@BrownianNotion I think the loss seems okay? The KL stuff I think we did change the way we computed the KL avg but not sure... main should be working okay these days!

Generally, I would just check to see if you see similar reward curves, and also just test downstream eval vs the released checkpoints.

[yukiwukii]

Hey @hamishivi , thanks so much! Any chance you know why the loss magnitudes are so different? Green is olmo, blue is my repro. Was the loss function, reduction method etc. changed? It's still a sum loss on finetune.py as far as I can see though.

Hey @BrownianNotion, I am trying to reproduce OLMo2-1B-SFT as well. My train loss curve is similar to yours. However, when I tested my trained model on downstream tasks, its performance is significantly worse than the original SFT model. Did you experience the same thing?

This is similar to issue #751, but in my custom eval, the difference between my trained model and the checkpoint is large (~20% point).

[BrownianNotion]

Hey @yukiwukii , I'm not sure what your custom eval is, but there was indeed a significant difference between the OLMo-2 1B SFT checkpoint on hugging face and my attempt at reproducing it from performing SFT on the base model. See my comment here #1352 (comment), there's about a 6% difference on GSM8K for example. If it's helpful to crosscheck, you can find my attempt at reconciling the reproduction commands here.