Now, dpo.py matches dpo_tune_cache.py almost perfectly on the single GPU experiments

[finbarrtimbers]

Colab with plots, particularly loss curve:

Changes:

1. In HFDataLoader, switched to match the dpo_tune_cache.py iteration order by using torch.randperm with seeded torch.Generator
2. Fix attention masking in the dpo.py forwards pass by converting padded batches to the packed format. We now pass doc_lens and max_doc_lens, enabling proper intra-document attention masking.
3. Fix weight loading order by loading after paralellization (this was previously re-initializing the weights, whoops).
4. Use 0-based epoch numbers to match dpo_tune_cache.py.
5. Use Duration.steps instead of Duration.epochs in dpo.py so we do max_train_steps training steps.

Changes not related to making DPO match:

1. Fixed the token_count metric in dpo_tune_cache.py to undo the mean reduction.
2. Make data ordering reproducible by using a seeded RandomSampler with torch.Generator().manual_seed(seed).

Known areas of deviation between the two:

1. dpo.py sets drop_last=True while dpo_tune_cache.py doesn't.
2. dpo.py accumulates gradients using a token-weighted mean, while dpo_tune_cache.py uses uniform weighting.

We think that dpo.py does things correctly.

---

Note

Medium Risk
Changes core DPO training semantics (data order, masking, checkpoint weight loading, and step scheduling), which can materially affect convergence and reproducibility; scope is contained to DPO paths and covered by updated GPU tests.

Overview
Makes dpo.py’s single-GPU behavior match dpo_tune_cache.py more closely by (1) switching HFDataLoader shuffling to seeded torch.randperm for reproducible ordering, (2) fixing OLMo-core attention masking by repacking padded batches into doc_lens/max_doc_lens format in dpo_utils forwards, and (3) reloading HF checkpoint weights after TransformerTrainModule parallelization to avoid accidental re-init.

Also updates training control to be step-based (max_train_steps via Duration.steps + scheduler step count), forces drop_last=True for DPO data/cache loaders, makes dpo_tune_cache.py sampling reproducible with a seeded RandomSampler, and corrects its token_count aggregation after mean reduction; tests and debug script are adjusted accordingly.

^{Written by Cursor Bugbot for commit c5cb1bb. This will update automatically on new commits. Configure here.}

[github-actions]

Documentation Changes Detected

📄 sitemap.xml

--- site-base/sitemap.xml	2026-02-18 19:40:01.684494357 +0000
+++ site-pr/sitemap.xml	2026-02-18 19:39:59.221033571 +0000
@@ -13,6 +13,10 @@
          <lastmod>2026-02-18</lastmod>
     </url>
     <url>
+         <loc>https://github.com/allenai/open-instruct/dpo_divergence_investigation/</loc>
+         <lastmod>2026-02-18</lastmod>
+    </url>
+    <url>

📄 sitemap.xml.gz

Binary files site-base/sitemap.xml.gz and site-pr/sitemap.xml.gz differ

Showing first 10 lines of diff for each changed file (up to 5 files, excluding search indices).

[github-actions]

Documentation Changes Detected

📄 sitemap.xml

--- site-base/sitemap.xml	2026-02-18 20:49:05.085290798 +0000
+++ site-pr/sitemap.xml	2026-02-18 20:49:02.708144399 +0000
@@ -13,6 +13,10 @@
          <lastmod>2026-02-18</lastmod>
     </url>
     <url>
+         <loc>https://github.com/allenai/open-instruct/dpo_divergence_investigation/</loc>
+         <lastmod>2026-02-18</lastmod>
+    </url>
+    <url>

📄 sitemap.xml.gz

Binary files site-base/sitemap.xml.gz and site-pr/sitemap.xml.gz differ

Showing first 10 lines of diff for each changed file (up to 5 files, excluding search indices).

[hamishivi]

generally lgtm but one comment

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

^{Bugbot Autofix is OFF. To automatically fix reported issues with Cloud Agents, enable Autofix in the Cursor dashboard.}

[cursor]

Cache data loader drops examples, causing RuntimeError

High Severity

The cache_data_loader was changed from drop_last=False to drop_last=True. This causes build_reference_logprobs_cache to skip caching some dataset indices (the remainder that doesn't fill a full batch). Since the cache function allocates tensors of size full_dataset_size=len(dataset) and then validates that every index was populated (raising RuntimeError for any -inf entries), this will crash whenever len(dataset) % cache_batch_size != 0.