Repository: huggingface/trl

Issue:
Negative KL warning even if KL='abs' and KL='full'
I'm using PPOTrainer to RLHF a llama2 model, and I sometimes get negative KL in my experiments. I've noticed that PPOTrainer prints a warning about that, even if I set `kl_penalty="abs"` or `kl_penalty="full"`, even though (by my understanding) it should be impossible to get negative KL in those cases.

Two questions:
* Checking the code, it looks like the KL divergence that's reported to wandb, and that is also used to determine whether the "negative KL" warning is printed, is calculated separately from the KL that is used for training; and the reported KL is always using the standard estimate. Did I understand that correctly, and does that mean I can ignore the "negative KL" warning if I'm using "abs" or "full" KL? (I can see that the non-score reward is non-positive if I use abs or full KL, and if I'm reading the code correctly the does indicate that the KL used in training is likely positive.)
* Conceptually, is there any issue or drawback with using "abs" or "full" instead of "kl"?

Thank you so much!


Discussion:
Indeed, that might be a bug, would you like to open a PR to fix it? In general `full` should be fine but we have observed some issues with the KL becoming negative and then `abs` is safer.
> Indeed, that might be a bug, would you like to open a PR to fix it? In general `full` should be fine but we have observed some issues with the KL becoming negative and then `abs` is safer.

Yes, happy to open a PR! Give me a few days.

And just to make sure my understanding of everything is correct: 
- With `kl` and negative KL, the real issue is that the negative KL leads to a positive non-score reward, and that opens up a loophole where the algorithm can get almost arbitrarily good returns just by driving down the KL into large-magnitude negative numbers.
- Negative KL with `kl` means that the log-prob of the chosen tokens were on average less than the log-probs of those same tokens in the reference model.
- That in itself (log-probs of chosen tokens being less than in the ref model) isn't necessarily an issue. It probably just means that the 'peak' of the token probability distribution decreased, i.e. the probability weight was spread out more across tokens, i.e. entropy increased. That could be indicative of an issue depending on what you're trying to do, but it doesn't *have* to be a problem. Is my intuition here roughly in line with how you are thinking about this?

And as an aside, can you get negative `full` KL? I thought KL by definition is non-negative; or is `full` still just an estimate pf the real KL?

Thank you so much!

Resolving patch:
diff --git a/trl/trainer/ppo_trainer.py b/trl/trainer/ppo_trainer.py
--- a/trl/trainer/ppo_trainer.py
+++ b/trl/trainer/ppo_trainer.py
@@ -733,11 +733,11 @@ def step(
                 active_full_logprobs = logprobs_from_logits(logits_or_none, None, gather=False)
                 ref_full_logprobs = logprobs_from_logits(ref_logits_or_none, None, gather=False)
 
-                rewards, non_score_reward = self.compute_rewards(
+                rewards, non_score_reward, kls = self.compute_rewards(
                     scores, active_full_logprobs, ref_full_logprobs, masks
                 )
             else:
-                rewards, non_score_reward = self.compute_rewards(scores, all_logprobs, ref_logprobs, masks)
+                rewards, non_score_reward, kls = self.compute_rewards(scores, all_logprobs, ref_logprobs, masks)
             timing["time/ppo/compute_rewards"] = time.time() - t
 
             t = time.time()
@@ -831,6 +831,7 @@ def step(
             masks=masks,
             queries=queries,
             responses=responses,
+            kls=kls,
         )
         # Gather/Reduce stats from all processes
         if self.is_distributed:
@@ -1091,11 +1092,17 @@ def compute_rewards(
                 Log probabilities of the model, shape (`batch_size`, `response_length`)
             ref_logprobs (`torch.FloatTensor`):
                 Log probabilities of the reference model, shape (`batch_size`, `response_length`)
+
+        Returns:
+            `torch.FloatTensor`: Per token rewards, shape (`batch_size`, `response_length`)
+            `torch.FloatTensor`: Non score rewards, shape (`batch_size`, `response_length`)
+            `torch.FloatTensor`: KL penalty, shape (`batch_size`, `response_length`)
         """
-        rewards, non_score_rewards = [], []
+        rewards, non_score_rewards, kls = [], [], []
         for score, logprob, ref_logprob, mask in zip(scores, logprobs, ref_logprobs, masks):
             # compute KL penalty (from difference in logprobs)
             kl = self._kl_penalty(logprob, ref_logprob)
+            kls.append(kl)
             non_score_reward = -self.kl_ctl.value * kl
             non_score_rewards.append(non_score_reward)
             reward = non_score_reward.clone()
@@ -1104,7 +1111,7 @@ def compute_rewards(
             # reward is preference model score + KL penalty
             reward[last_non_masked_index] += score
             rewards.append(reward)
-        return torch.stack(rewards), torch.stack(non_score_rewards)
+        return torch.stack(rewards), torch.stack(non_score_rewards), torch.stack(kls)
 
     def _kl_penalty(self, logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor) -> torch.FloatTensor:
         if self.config.kl_penalty == "kl":
@@ -1256,7 +1263,8 @@ def record_step_stats(self, kl_coef: float, **data):
         """
         mask = data.pop("masks")
 
-        kl_list = ((data["logprobs"] - data["ref_logprobs"]) * mask).sum(axis=-1)
+        kls = data.pop("kls")
+        kl_list = ((kls) * mask).sum(axis=-1)
         mean_kl = kl_list.mean()
         mean_entropy = (-data["logprobs"] * mask).sum(axis=-1).mean()
 


Test patch:
diff --git a/tests/test_ppo_trainer.py b/tests/test_ppo_trainer.py
--- a/tests/test_ppo_trainer.py
+++ b/tests/test_ppo_trainer.py
@@ -579,7 +579,7 @@ def test_loss_trainer(self):
         logits = torch.exp(all_logprobs)
         vpreds = values + 0.1
 
-        score, non_score = ppo_trainer.compute_rewards(dummy_scores, all_logprobs, ref_logprobs, mask)
+        score, non_score, kls = ppo_trainer.compute_rewards(dummy_scores, all_logprobs, ref_logprobs, mask)
         values, advantages, returns = ppo_trainer.compute_advantages(values, score, mask)
 
         # just make sure a dummy loss is computed