Support for pairwise judges in online training

src/fairseq2/recipes/lm/__init__.py

@@ -171,11 +171,17 @@
 from fairseq2.recipes.lm._online_finetune._rewards import (
     GenerativePointwiseVerifier as GenerativePointwiseVerifier,
 )
-
 from fairseq2.recipes.lm._online_finetune._rewards import (
     GenerativePointwiseVerifierHandler as GenerativePointwiseVerifierHandler,
 )
 
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    GenerativePairwiseVerifier as GenerativePairwiseVerifier,
+)
+from fairseq2.recipes.lm._online_finetune._rewards import (
+    GenerativePairwiseVerifierHandler as GenerativePairwiseVerifierHandler,
+)
+
 from fairseq2.recipes.lm._online_finetune._remote_model import (
     RemoteModelHandler as RemoteModelHandler,
 )

src/fairseq2/recipes/lm/_online_finetune/_common.py

@@ -389,7 +389,15 @@ def generate_rewards(
     return rewards_per_rank[0]
 
 
-def generate_rewards_generative(prompts: List[List[int]], dp_gang, vllm_model):
+def generate_rewards_generative(
+    prompts: List[List[int]], 
+    dp_gang, 
+    vllm_model,
+    is_pointwise=True
+):
+    """
+    By default, any LLM-as-a-Judge will be used in a pointwise setup that generates a score
+    """
     prompts_to_generate = [None] * dp_gang.size
     if dp_gang.rank == 0:
         dp_gang.gather_object(prompts, prompts_to_generate, 0)
@@ -401,7 +409,7 @@ def generate_rewards_generative(prompts: List[List[int]], dp_gang, vllm_model):
         for rank_prompts in prompts_to_generate:
             flat_request_list.extend(rank_prompts)
 
-        rewards = vllm_model.reward_from_generative_model(flat_request_list)
+        rewards = vllm_model.reward_from_generative_model(flat_request_list, is_pointwise)
 
         rewards_to_scatter = []
         rewards_per_rank = [None]
@@ -642,6 +650,8 @@ def log_rollouts(prompt_batch: PromptBatch, rollouts, split_name, num_rollouts=1
     """
     if "prompt_raw" in prompt_batch.meta_info:
         prompt = prompt_batch.meta_info.get("prompt_raw")[0]
+    elif "raw_prompt" in prompt_batch.meta_info:
+        prompt = prompt_batch.meta_info.get("raw_prompt")[0]
     else:
         # raw text prompt doesn't exist for this dataset
         prompt = "DUMMY PROMPT"

src/fairseq2/recipes/lm/_online_finetune/_generative_prompts.py

@@ -18,3 +18,55 @@
 {response}
 [The End of the Assistant's Answer]
 """
+
+PAIRWISE_PROMPT = """
+You are given a user question and two responses from two AI assistants. Your task is to act as an impartial judge and evaluate which response better follows the user's instructions and provides a higher-quality answer.
+
+First, provide your reasoning within <think> and </think> tags. This should include your evaluation criteria for a high-quality response, a detailed comparison of the two responses, and when helpful, a reference answer as part of your evaluation. Be explicit in your thought process, referencing your criteria and explaining how each response aligns with or deviates from them.
+
+Avoid any position biases and ensure that the order in which the responses were presented does not influence your decision. Do not allow the length of the responses to influence your evaluation. Do not favor certain names of the assistants. Be as objective as possible.
+
+Finally, provide your verdict within <answer> and </answer> tags, strictly following this format:
+- <answer> [[A]] </answer> if Assistant A is better
+- <answer> [[B]] </answer> if Assistant B is better
+
+Below are the user's question and the two responses:
+
+[User Question]
+{instruction}
+
+[The Start of Assistant A's Answer]
+{response_A}
+[The End of Assistant A's Answer]
+
+[The Start of Assistant B's Answer]
+{response_B}
+[The End of Assistant B's Answer]
+"""
+
+PAIRWISE_WITH_SCORES_PROMPT = """
+You are given a user question and two responses from two AI assistants. Your task is to act as an impartial judge and evaluate which response better follows the user's instructions and provides a higher-quality answer.
+
+First, provide your reasoning within <think> and </think> tags. This should include your evaluation criteria for a high-quality response, a detailed comparison of the two responses, and when helpful, a reference answer as part of your evaluation. Be explicit in your thought process, referencing your criteria and explaining how each response aligns with or deviates from them.
+
+Avoid any position biases and ensure that the order in which the responses were presented does not influence your decision. Do not allow the length of the responses to influence your evaluation. Do not favor certain names of the assistants. Be as objective as possible.
+
+Finally, assign the assistant's response a score from 0 to 10, using either an integer or a decimal with up to 0.1 precision, with a higher score indicating a higher-quality response that better satisfies the criteria. Enclose the scores within the tags <score_A> </score_A>, and <score_B> </score_B>.
+
+Format your output like this:
+<think> your_thinking_process </think>
+<score_A> your_score_a </score_A> <score_B> your_score_b </score_B>
+
+Below are the user's question and the two responses:
+
+[User Question]
+{instruction}
+
+[The Start of Assistant A's Answer]
+{response_A}
+[The End of Assistant A's Answer]
+
+[The Start of Assistant B's Answer]
+{response_B}
+[The End of Assistant B's Answer]
+"""

src/fairseq2/recipes/lm/_online_finetune/_grpo.py

@@ -75,7 +75,6 @@
 from fairseq2.utils.structured import structure
 from fairseq2.utils.validation import validate
 
-
 @final
 class GrpoFinetuneUnit(TrainUnit[SequenceBatch]):
     """Represents the language model DPO-finetuning unit with online generations. Paper: https://arxiv.org/abs/2305.18290."""
@@ -89,6 +88,7 @@ class GrpoFinetuneUnit(TrainUnit[SequenceBatch]):
     _sync_vllm_model_every_n_steps: int
     _sync_ref_model_every_n_steps: int
     _reward: VLLMOutputReward
+    _reward_name: str
     _display_name: str
     _reference_offload: bool
     _rollout_bag: StatefulRolloutBag
@@ -101,6 +101,7 @@ def __init__(
         vllm_model: RemoteVllmModel,
         vllm_actors: List[Union[RemoteVllmModel, RemoteHFModel]],
         reward,
+        reward_name: str,
         gangs: Gangs,
         loss_config: GrpoLossConfig,
         sync_vllm_model_every_n_steps: int = 1,
@@ -116,6 +117,7 @@ def __init__(
         self._sync_vllm_model_every_n_steps = sync_vllm_model_every_n_steps
         self._sync_ref_model_every_n_steps = sync_ref_model_every_n_step
         self._reward = reward
+        self._reward_name = reward_name
         self._reference_offload = reference_offload
         self._metric_bag = GrpoFinetuneMetricBag(gangs.dp)
         self._rollout_bag = StatefulRolloutBag()
@@ -161,7 +163,8 @@ def maybe_sync_models(self, force_sync_vllm=False):
     def validate_reward(self, prompt_batch: PromptBatch) -> tuple[Tensor, int]:
         if self._gangs.dp.rank == 0:
             policy_sampling_params = copy(self._vllm_model.sampling_params)
-            policy_sampling_params.n = 1
+            # For a pairwise RM, need to sample at least two judgments
+            policy_sampling_params.n = 2 if self._reward_name == "generative_pairwise_verifier" else 1
             for k, v in self._loss_config.validation_vllm_sampling_params.items():
                 policy_sampling_params.__setattr__(k, v)
         else:
@@ -175,7 +178,9 @@ def validate_reward(self, prompt_batch: PromptBatch) -> tuple[Tensor, int]:
         if self._loss_config.log_rollouts:
             log_rollouts(prompt_batch, rollouts, "Valid")
         reward_output = self._reward.process_rollouts(rollouts, prompt_batch)
+        log.info(f"Rewards: {reward_output['rewards']}")
         avg_reward = torch.tensor(reward_output["rewards"]).float().mean()
+        std_reward = torch.tensor(reward_output["rewards"]).float().std()
 
         rollout_lengths = get_rollout_lengths(rollouts)
         avg_rollout_length = torch.tensor(rollout_lengths).float().mean()
@@ -185,6 +190,7 @@ def validate_reward(self, prompt_batch: PromptBatch) -> tuple[Tensor, int]:
         self._metric_bag.update_avg_reward_len_norm(avg_reward_len_norm)
 
         self._metric_bag.update_avg_reward(avg_reward)
+        self._metric_bag.update_std_reward(std_reward)
         self._metric_bag.update_batch_metrics(prompt_batch)
         # returning dummy loss since trainer expects it
         return torch.tensor(0.0, device=self._gangs.dp.device), prompt_batch.batch_size
@@ -314,7 +320,10 @@ def __call__(self, prompt_batch: PromptBatch) -> tuple[Tensor, int]:
         )  # TODO fix, now logs only the last prompt from the batch
 
         avg_reward = torch.tensor(reward_output["rewards"]).float().mean()
+        std_reward = torch.tensor(reward_output["rewards"]).float().std()
+
         self._metric_bag.update_avg_reward(avg_reward)
+        self._metric_bag.update_std_reward(std_reward)
 
         loss = grpo_loss
 
@@ -400,6 +409,7 @@ class GrpoFinetuneMetricBag(SequenceMetricBag):
     grpo_loss: Mean
     logit_entropy: Mean
     avg_reward: Mean
+    std_reward: Mean
 
     def __init__(self, gang: Gang) -> None:
         super().__init__(gang)
@@ -419,6 +429,9 @@ def __init__(self, gang: Gang) -> None:
         self.register_metric(
             "logit_entropy", Mean(device=gang.device), persistent=False
         )
+        self.register_metric(
+            "std_reward", Mean(device=gang.device), persistent=False
+        )
 
     @torch.inference_mode()
     def update_logit_entropy(self, logit_entropy: Tensor):
@@ -455,6 +468,10 @@ def update_rollout_lengths(
     @torch.inference_mode()
     def update_avg_reward(self, avg_reward):
         self.avg_reward.update(avg_reward, weight=1)
+
+    @torch.inference_mode()
+    def update_std_reward(self, std_reward):
+        self.std_reward.update(std_reward, weight=1)
 
     @torch.inference_mode()
     def update_avg_rollout_length(self, avg_rollout_length):
@@ -617,7 +634,8 @@ def create(
 
         vllm_reward_model = vllm_actors.get(config.vllm_reward_model_name, None)
         reward_registry = self._context.get_registry(VLLMOutputRewardHandler)
-        reward_handler = reward_registry.get(config.reward.name)
+        reward_name = config.reward.name
+        reward_handler = reward_registry.get(reward_name)
         reward = reward_handler.create(
             reward_model=vllm_reward_model,
             reward_config=config.reward.config,
@@ -633,6 +651,7 @@ def create(
             vllm_model,
             vllm_actors,
             reward,
+            reward_name,
             gangs,
             config.loss_config,
             config.sync_vllm_model_every_n_steps,

src/fairseq2/recipes/lm/_online_finetune/_remote_model.py

@@ -7,6 +7,7 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections import Counter
 from dataclasses import dataclass, field
 from typing_extensions import override
 from vllm.engine.arg_utils import PoolerConfig
@@ -263,42 +264,76 @@ def reward_from_model(self, prompt_list, batch_size=64):
         rewards = [o.outputs.data.item() for o in ray_outputs_flat]
         return rewards
 
-    def reward_from_generative_model(self, prompt_list):
+    def reward_from_generative_model(self, prompt_list, is_pointwise):
 
-        def extract_score(output):
+        def extract_score_single(output):
             matches = re.findall(
                 r"<score>\s*([0-9]+(?:\.[0-9])?)\s*(?:/10)?\s*</score>", output
             )
-            return float(matches[-1]) if matches else 0.0
-
+            if matches and float(matches[-1].strip()) > 10.0:
+                log.info(f"CoT = {output}")
+            return float(matches[-1].strip()) if matches else 0.0
+
+        def extract_score_pair(output):
+            score_a_matches = re.findall(r"<score_A>\s*([0-9]+(?:\.[0-9])?)\s*(?:/10)?\s*</score_A>", output)
+            score_b_matches = re.findall(r"<score_B>\s*([0-9]+(?:\.[0-9])?)\s*(?:/10)?\s*</score_B>", output)
+
+            if score_a_matches and score_b_matches:
+                score_a = score_a_matches[-1]  # Last occurrence of score_A
+                score_b = score_b_matches[-1]  # Last occurrence of score_B
+                if float(score_a.strip()) > 10.0 or float(score_b.strip()) > 10.0:
+                    log.info(f"CoT = {output}")
+                return (float(score_a.strip()), float(score_b.strip()))
+            else:
+                return (0.0, 0.0)
+
+        def get_avg_score(scores, is_pointwise):
+            avg_score = 0.0 if is_pointwise else (0.0, 0.0)
+            for score in scores:
+                if is_pointwise:
+                    avg_score += score
+                else:
+                    avg_score = (avg_score[0]+score[0], avg_score[1]+score[1])
+
+            if is_pointwise:
+                return round(avg_score/len(scores), 4)
+            else:
+                return (round(avg_score[0]/len(scores), 4), round(avg_score[1]/len(scores), 4))
+
         def get_len_norm_avg_score(scores, lengths):
             avg_score = 0.0
             for score, length in zip(scores, lengths):
-                avg_score += score / length
-
-            return round(avg_score / len(scores), 4)
-
-        def get_avg_score(scores):
-            avg_score = 0.0
-            for score in scores:
-                avg_score += score
-
-            return round(avg_score / len(scores), 4)
+                avg_score += score/length
+
+            return round(avg_score/len(scores), 4)
+
+        def extract_preference(output):
+            matches = list(
+                re.finditer(r"<answer>\s*\[\[(A|B)\]\]\s*</answer>", output.strip())
+            )
 
-        rewards = []
+            return matches[-1].group(1) if matches else None
+
+        def get_majority_vote(preferences):
+            count_A = preferences.count("A")
+            count_B = preferences.count("B")
+            if count_A > count_B:
+                return "A"
+            elif count_A < count_B:
+                return "B"
+            else:
+                return None
+
         judgments = self.rollout_from_model(prompt_list=prompt_list, string_input=True)
-
+        
         rewards = []
         for per_rollout_judgments in judgments:
-            per_rollout_scores = [
-                extract_score(judgment.text)
-                for judgment in per_rollout_judgments.outputs
-            ]
-            per_rollout_lengths = [
-                len(judgment.token_ids) for judgment in per_rollout_judgments.outputs
-            ]
-            rewards.append(get_avg_score(per_rollout_scores))
-
+            if is_pointwise:
+                per_rollout_scores = [extract_score_single(judgment.text) for judgment in per_rollout_judgments.outputs]
+            else:
+                per_rollout_scores = [extract_score_pair(judgment.text) for judgment in per_rollout_judgments.outputs]
+            rewards.append(get_avg_score(per_rollout_scores, is_pointwise))
+
         return rewards
 
 
@@ -485,4 +520,4 @@ def name(self) -> str:
     @property
     @override
     def config_kls(self) -> type[object]:
-        return RayActorConfig
+        return RayActorConfig