Jacklanchantin/verifiable nonverifiable

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

@@ -62,6 +62,7 @@
     RewardSection,
     VLLMOutputReward,
     VLLMOutputRewardHandler,
+    MultiVerifier,
 )
 from fairseq2.recipes.lm._preference_finetune._common import (
     POCriterionSection,
@@ -89,7 +90,7 @@ class OnlineDpoFinetuneUnit(TrainUnit[SequenceBatch]):
     _sync_vllm_model_every_n_steps: int
     _sync_ref_model_every_n_steps: int
     _display_name: str
-    _reward: VLLMOutputReward
+    _reward: dict | VLLMOutputReward
     _reference_offload: bool
 
     def __init__(
@@ -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 = reward
         self._metric_bag = OnlineDpoFinetuneMetricBag(gangs.dp)
 
         self._display_name = "online_dpo"
@@ -225,9 +227,17 @@ def __call__(self, prompt_batch: PromptBatch) -> tuple[Tensor, int]:
             log_rollouts(prompt_batch, rollouts, "Train")
 
         batch: PreferenceBatch
+        # Class MultiReward (superclass of current reward)
+        # subrewards AtheneReward, MathReward
+        # this super class should be working as is for 2 validation sets, since validation sets will also have task type in the jsonl, so hypothetically correct reward will also be used for logging
+        # if self._gangs.root.rank == 0:
+        #     breakpoint()
+        # self._gangs.root.barrier()
+
         batch, is_bad_batch, reward_output = self._reward.prepare_preference_batch(
             prompt_batch, rollouts
         )  # loss_zeroer is used when entire batch has no valid prefrence pair
+
         if is_bad_batch:
             loss_zeroer = 0.0
         else:
@@ -562,15 +572,26 @@ def create(
         gangs.root.barrier()
 
         vllm_model = vllm_actors[config.ray_policy_actor_name]
-
-        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 = reward_handler.create(
-            reward_model=vllm_reward_model,
-            reward_config=config.reward.config,
-            gangs=gangs,
-        )
+        vllm_reward_model = vllm_actors.get(config.vllm_reward_model_name, None)
+        if type(config.reward.name) is list:
+            reward_mapper = {}
+            for reward_name in config.reward.name:
+                reward_handler = reward_registry.get(reward_name)
+                reward_model = reward_handler.create(
+                    reward_model=vllm_reward_model,
+                    reward_config=config.reward.config,
+                    gangs=gangs,
+                )
+                reward_mapper[reward_name] = reward_model
+            reward = MultiVerifier(reward_mapper)
+        else:
+            reward_handler = reward_registry.get(reward_name)
+            reward = reward_handler.create(
+                reward_model=vllm_reward_model,
+                reward_config=config.reward.config,
+                gangs=gangs,
+            )
 
         return OnlineDpoFinetuneUnit(
             model,

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

@@ -9,7 +9,7 @@
 import re
 from abc import ABC, abstractmethod
 from dataclasses import dataclass, field
-from typing import Any, List
+from typing import Any, List, Dict
 
 import torch
 from transformers import AutoTokenizer
@@ -34,6 +34,7 @@
 )
 from fairseq2.recipes.model import Model
 from fairseq2.recipes.trainer import TrainUnit
+from fairseq2.logging import log
 
 
 @dataclass(kw_only=True)
@@ -45,7 +46,7 @@ class RewardModelConfig:
 
 @dataclass(kw_only=True)
 class RewardSection:
-    name: str = "dummy"
+    name: str | List[str] = "dummy"
     config: RewardModelConfig = field(default_factory=lambda: RewardModelConfig())
 
 
@@ -179,6 +180,7 @@ def create(self, reward_model, reward_config, gangs):
             answer_key=reward_config.answer_key,
             prompt_key=reward_config.prompt_key,
             gangs=gangs,
+            vllm_math_reward_model=reward_model,
         )
 
     @property
@@ -193,7 +195,7 @@ def config_kls(self):
 
 
 class MathVerifyVerifier(VLLMOutputReward):
-    def __init__(self, answer_key, prompt_key, gangs):
+    def __init__(self, answer_key, prompt_key, gangs, vllm_math_reward_model):
         try:
             from math_verify.metric import math_metric
             from math_verify.parser import (
@@ -209,6 +211,8 @@ def __init__(self, answer_key, prompt_key, gangs):
         self._gangs = gangs
         self.answer_key = answer_key
         self.prompt_key = prompt_key
+        self.vllm_math_reward_model = vllm_math_reward_model
+        self.tokenizer = AutoTokenizer.from_pretrained("Nexusflow/Athene-RM-8B")
 
         label_normalizer = NormalizationConfig(
             basic_latex=True,
@@ -227,6 +231,16 @@ def __init__(self, answer_key, prompt_key, gangs):
             precision=6,
         )
 
+    def wrap_text(self, prompt_text, rollout_text):
+        wrapped_text = [
+            {"role": "user", "content": prompt_text},
+            {"role": "assistant", "content": rollout_text},
+        ]
+        chat_str = self.tokenizer.apply_chat_template(wrapped_text, tokenize=False)
+        chat_str += "<|reserved_special_token_1|>"
+
+        return chat_str
+
     def verify_answer(self, completion: str, answer: str):
         # here we add extra $$ to label so that LatexExtractor works as expected
         if not answer.startswith("$"):
@@ -246,12 +260,13 @@ def process_rollouts(
         vllm_outputs: List[RequestOutput],
         prompt_batch: PromptBatch,
     ):
+        log.info("MathVerifyVerifier process_rollouts()")
         batch_text = []
         batch_tokens = []
         batch_rewards = []
+        vllm_inputs = []  # dummy for vllm syncronization # FIXME
 
         reference_answers = prompt_batch.meta_info.get(self.answer_key)
-
         for i, i_batch_request_output in enumerate(vllm_outputs):
             rollouts_text = []
             rollouts_tokens = []
@@ -264,16 +279,29 @@ def process_rollouts(
                     rollout_output.text, i_reference_answer
                 )
                 rollouts_rewards.append(predicted_reward)
+                vllm_input = self.wrap_text("dummy prompt", "dummy rollout")
+                vllm_inputs.append(vllm_input)
+
             batch_text.append(rollouts_text)
             batch_tokens.append(rollouts_tokens)
             batch_rewards.append(rollouts_rewards)
 
+        # dummy vllm call for syncronization # FIXME
+        # if self.vllm_math_reward_model is not None:
+        log.info("MathVerifyVerifir generate_rewards()")
+        dummy_rewards = generate_rewards(
+            vllm_inputs,
+            dp_gang=self._gangs.dp,
+            vllm_model=self.vllm_math_reward_model,
+        )
+
         return {"text": batch_text, "tokens": batch_tokens, "rewards": batch_rewards}
 
     def prepare_preference_batch(
         self, prompt_batch: PromptBatch, rollouts
     ) -> PreferenceBatch:
 
+        log.info("MathVerifyVerifier prepare_preference_batch()")
         reward_output = self.process_rollouts(rollouts, prompt_batch)
 
         batch, is_bad_batch = prepare_preference_batch_random_pair(
@@ -601,6 +629,7 @@ def wrap_text(self, prompt_text, rollout_text):
     def process_rollouts(
         self, vllm_outputs: List[RequestOutput], prompt_batch: PromptBatch
     ):
+        log.info("AtheneVerifier process_rollouts()")
         vllm_inputs = []
         batch_text = []
         batch_tokens = []
@@ -625,6 +654,7 @@ def process_rollouts(
             batch_text.append(rollouts_text)
             batch_tokens.append(rollouts_tokens)
 
+        log.info("AtheneVerifier generate_rewards()")
         batch_rewards = generate_rewards(
             vllm_inputs, dp_gang=self._gangs.dp, vllm_model=self.reward_model
         )
@@ -639,6 +669,7 @@ def prepare_preference_batch(
         self, prompt_batch: PromptBatch, rollouts
     ) -> PreferenceBatch:
 
+        log.info("AtheneVerifier prepare_preference_batch()")
         reward_output = self.process_rollouts(rollouts, prompt_batch)
 
         chosen_batch = []
@@ -751,3 +782,52 @@ def prepare_grpo_batch(self, prompt_batch: PromptBatch, rollouts):
         )
 
         return grpo_batch, reward_output
+
+
+class MultiVerifier(VLLMOutputReward):
+    """
+    A reward model verifier that processes rollouts using multiple reward models.
+    This class evaluates rollouts generated by vLLM by wrapping the prompt and rollout text into a specific format and passing it through the appropriate reward model.
+    This class is designed to handle multiple reward models and their configurations.
+    Note: this relies on modified Athene-RM-8B code to ensure compatibility with vLLM.
+    """
+
+    def __init__(self, rewards_map: Dict[VLLMOutputReward]):
+        self.rewards_map = rewards_map
+
+    def get_reward_model_type(self, prompt_batch: PromptBatch) -> str:
+        """
+        Get the reward model type from the prompt batch meta info.
+        """
+        assert (
+            len(set(prompt_batch.meta_info["reward_model"])) == 1
+        ), "Not all batch prompts have the same 'reward_model' value."
+        reward_model_type = prompt_batch.meta_info["reward_model"][0]
+        if reward_model_type is None:
+            raise ValueError("Reward model type not found in meta_info.")
+
+        log.info(f"Reward model type: {reward_model_type}")
+        return reward_model_type
+
+    @override
+    def process_rollouts(
+        self, vllm_outputs: List[RequestOutput], prompt_batch: PromptBatch
+    ):
+        reward_model_type = self.get_reward_model_type(prompt_batch)
+        reward = self.rewards_map[reward_model_type]
+        return reward.process_rollouts(vllm_outputs, prompt_batch)
+
+    @override
+    def prepare_preference_batch(
+        self, prompt_batch: PromptBatch, rollouts
+    ) -> PreferenceBatch:
+
+        reward_model_type = self.get_reward_model_type(prompt_batch)
+        reward = self.rewards_map[reward_model_type]
+        return reward.prepare_preference_batch(prompt_batch, rollouts)
+
+    @override
+    def prepare_grpo_batch(self, prompt_batch: PromptBatch, rollouts):
+        reward_model_type = self.get_reward_model_type(prompt_batch)
+        reward = self.rewards_map[reward_model_type]
+        return reward.prepare_grpo_batch(prompt_batch, rollouts)