Log chosen/rejected entropy

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

@@ -343,9 +343,7 @@ def _compute_grpo_objective(
                 (per_token_loss * target_mask).sum(dim=-1) / target_mask.sum(dim=-1)
             ).mean(dim=1)
         else:
-            per_seq_loss = (
-                (per_token_loss * target_mask).sum(dim=-1)
-            ).mean(dim=1)
+            per_seq_loss = ((per_token_loss * target_mask).sum(dim=-1)).mean(dim=1)
 
         # if self._gangs.root.rank == 0:
         #     from pudb.remote import set_trace

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

@@ -208,6 +208,16 @@ def compute_reference_logps(self, seq_batch: SequenceBatch):
 
         return ref_logps
 
+    def get_all_rollouts_entropy(self, rollouts):
+        all_entropy = []
+        for rollout_idx in range(len(rollouts[0].outputs)):
+            logprobs = rollouts[0].outputs[rollout_idx].logprobs
+            logprobs = [next(iter(x.values())).logprob for x in logprobs]
+            entropy = sum(logprobs) / len(logprobs)
+            all_entropy.append(entropy)
+        logit_entropy = torch.tensor(all_entropy, device=self._gangs.dp.device)
+        return logit_entropy
+
     @override
     def __call__(self, prompt_batch: PromptBatch) -> tuple[Tensor, int]:
 
@@ -263,14 +273,24 @@ def __call__(self, prompt_batch: PromptBatch) -> tuple[Tensor, int]:
         rejected_logps, average_rejected_logps = _gather_lprobs_avg(
             rejected_output, rejected_target_batch
         )
-        tgt_logit_entropy = compute_token_level_entropy(
+        chosen_tgt_logit_entropy = compute_token_level_entropy(
             chosen_output.logits, chosen_target_batch.target_mask
         )  # [Batch x Rollouts, 1]
+        rejected_tgt_logit_entropy = compute_token_level_entropy(
+            rejected_output.logits, rejected_target_batch.target_mask
+        )  # [Batch x Rollouts, 1]
+
+        # entropy for all N rollouts
+        logit_entropy = self.get_all_rollouts_entropy(rollouts)
 
         max_entropy_regularizer = (
-            -tgt_logit_entropy.sum() * self._loss_config.entropy_regularizer_scale
+            -chosen_tgt_logit_entropy.sum()
+            * self._loss_config.entropy_regularizer_scale
         )
-        self.metric_bag.update_logit_entropy(tgt_logit_entropy)
+
+        self.metric_bag.update_logit_entropy(logit_entropy)
+        self.metric_bag.update_chosen_logit_entropy(chosen_tgt_logit_entropy)
+        self.metric_bag.update_rejected_logit_entropy(rejected_tgt_logit_entropy)
 
         if self._reference_offload:
             token_ref_chosen_logps = self.compute_reference_logps(batch.chosen)
@@ -399,6 +419,8 @@ class OnlineDpoFinetuneMetricBag(POFinetuneMetricBag):
     num_dummy_batches: Mean
     avg_reward: Mean
     avg_loss_zeroer: Mean
+    chosen_logit_entropy: Mean
+    rejected_logit_entropy: Mean
     logit_entropy: Mean
 
     def __init__(self, gang: Gang) -> None:
@@ -416,6 +438,22 @@ def __init__(self, gang: Gang) -> None:
             "logit_entropy", Mean(device=gang.device), persistent=False
         )
 
+    @torch.inference_mode()
+    def update_chosen_logit_entropy(self, logit_entropy: Tensor):
+        # logit_entropy for chosen sequences
+        batch_size = logit_entropy.size(0)
+        self.chosen_logit_entropy.update(
+            logit_entropy.sum() / batch_size, weight=batch_size
+        )
+
+    @torch.inference_mode()
+    def update_rejected_logit_entropy(self, logit_entropy: Tensor):
+        # logit_entropy for rejected sequences
+        batch_size = logit_entropy.size(0)
+        self.rejected_logit_entropy.update(
+            logit_entropy.sum() / batch_size, weight=batch_size
+        )
+
     @torch.inference_mode()
     def update_logit_entropy(self, logit_entropy: Tensor):
         # logit_entropy is expected to contain token-level entropy for every sequence in the current batch

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

@@ -221,7 +221,7 @@ def rollout_from_model(self, prompt_list, sampling_params=None):
 
         return outputs
 
-    def reward_from_model(self, prompt_list, batch_size=64):
+    def reward_from_model(self, prompt_list, batch_size=16):
         # NOTE: need to batch inputs to vllm.encode model for current models that aren't supported by vllm
         rewards = []
         for i in range(0, len(prompt_list), batch_size):

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

@@ -365,22 +365,11 @@ def __init__(self, gangs, reward_model, answer_key, prompt_key):
             "Skywork/Skywork-Reward-Llama-3.1-8B-v0.2"
         )
 
-    def extract_text_from_llama3_wrapper(self, input_string):
-        start_pattern = r"<\|start_header_id\|>user<\|end_header_id\|>"
-        end_pattern = r"<\|eot_id\|><\|start_header_id\|>assistant<\|end_header_id\|>"
-        start_index = re.search(start_pattern, input_string).end()
-        end_index = re.search(end_pattern, input_string).start()
-        # Extract the text between the start and end indices
-        extracted_text = input_string[start_index:end_index].strip()
-        return extracted_text
-
     def wrap_text(self, prompt_text, rollout_text):
         wrapped_text = [
             {"role": "user", "content": prompt_text},
             {"role": "assistant", "content": rollout_text},
         ]
-        # templated_text = self.tokenizer.apply_chat_template(wrapped_text, tokenize=True)
-        # tokens_prompt = TokensPrompt(prompt_token_ids=templated_text)
         chat_str = self.tokenizer.apply_chat_template(wrapped_text, tokenize=False)
         chat_str = chat_str.replace("<|begin_of_text|>", "")
 

src/fairseq2/setup/_metrics.py

@@ -69,18 +69,20 @@ def register(name: str, *args: Any) -> None:
     register("generator_cache_capacity",      "Generator/Cache Capacity",      904, format_as_byte_size)
 
     # Preference Optimization
-    register("cpo_loss",         "CPO Loss",                             0, format_as_float)
-    register("dpo_loss",         "DPO Loss",                             0, format_as_float)
-    register("orpo_loss",        "ORPO Loss",                            0, format_as_float)
-    register("simpo_loss",       "SimPO Loss",                           0, format_as_float)
-    register("grpo_loss",        "GRPO Loss",                            0, format_as_float)
-    register("avg_reward",       "Reward",                               1, format_as_float)
-    register("chosen_logps",     "Chosen Sequence Log Probabilities",   50, format_as_float)
-    register("rejected_logps",   "Rejected Sequence Log Probabilities", 50, format_as_float)
-    register("logit_entropy",    "Logit Entropy",                       51, format_as_float)
-    register("rollout_lengths",  "Rollout Length",                      70, format_as_float)
-    register("chosen_lengths",   "Chosen Sequence Length",              70, format_as_float)
-    register("rejected_lengths", "Rejected Sequence Length",            70, format_as_float)
+    register("cpo_loss",              "CPO Loss",                             0, format_as_float)
+    register("dpo_loss",              "DPO Loss",                             0, format_as_float)
+    register("orpo_loss",             "ORPO Loss",                            0, format_as_float)
+    register("simpo_loss",            "SimPO Loss",                           0, format_as_float)
+    register("grpo_loss",             "GRPO Loss",                            0, format_as_float)
+    register("avg_reward",            "Reward",                               1, format_as_float)
+    register("chosen_logps",          "Chosen Sequence Log Probabilities",   50, format_as_float)
+    register("rejected_logps",        "Rejected Sequence Log Probabilities", 50, format_as_float)
+    register("logit_entropy",         "Logit Entropy",                       51, format_as_float)
+    register("chosen_logit_entropy",  "Chosen Logit Entropy",                51, format_as_float)
+    register("rejected_logit_entropy","Rejected Logit Entropy",              51, format_as_float)
+    register("rollout_lengths",       "Rollout Length",                      70, format_as_float)
+    register("chosen_lengths",        "Chosen Sequence Length",              70, format_as_float)
+    register("rejected_lengths",      "Rejected Sequence Length",            70, format_as_float)
 
     # Memory
     register("peak_active_mem",         "Peak Active Device Memory",       920, format_as_byte_size)