From 8c6e94e779148981530c7307406d9f4f326c12c0 Mon Sep 17 00:00:00 2001
From: Ilia Kulikov <kulikov@meta.com>
Date: Sat, 20 Sep 2025 02:49:10 +0000
Subject: [PATCH] adding prompt lengths to athene reward so that online dpo
 works as expected

---
 src/fairseq2/recipes/lm/_online_finetune/_recipe.py | 13 ++++++++-----
 .../recipes/lm/_online_finetune/_rewards.py         |  7 +++++++
 2 files changed, 15 insertions(+), 5 deletions(-)

diff --git a/src/fairseq2/recipes/lm/_online_finetune/_recipe.py b/src/fairseq2/recipes/lm/_online_finetune/_recipe.py
index 1317bf512..67d0e7884 100644
--- a/src/fairseq2/recipes/lm/_online_finetune/_recipe.py
+++ b/src/fairseq2/recipes/lm/_online_finetune/_recipe.py
@@ -252,11 +252,14 @@ def load_online_finetuner(
 
     vocab_size = tokenizer.vocab_info.size
 
-    # initialize ray and vllm actors
-    ray.init(
-        address=f"ray://{config.vllm.ray_cluster_ip_address}:10001",
-        namespace="vllm_workers",
-    )
+    if gangs.root.rank == 0:
+        # initialize ray and vllm actors
+        ray.init(
+            address=f"ray://{config.vllm.ray_cluster_ip_address}:10001",
+            namespace="vllm_workers",
+        )
+    
+    gangs.root.barrier()
 
     vllm_actors = {}
     # go over actor configs and initialize all of them
diff --git a/src/fairseq2/recipes/lm/_online_finetune/_rewards.py b/src/fairseq2/recipes/lm/_online_finetune/_rewards.py
index 49c7d9bcd..b9048cdd1 100644
--- a/src/fairseq2/recipes/lm/_online_finetune/_rewards.py
+++ b/src/fairseq2/recipes/lm/_online_finetune/_rewards.py
@@ -477,6 +477,13 @@ def prepare_preference_batch(
             reference_score_rejected=None,
         )
 
+        prompt_lengths = [
+            l
+            for idx, l in enumerate(prompt_batch.prompt_lengths)
+            if idx not in dummy_batch_ids
+        ]
+        reward_output["prompt_lengths"] = prompt_lengths
+
         return batch, is_bad_batch, reward_output