[BugFix] Normalize reward loss over valid pairs (#3886)

Author: fallintoplace

test/test_rlhf.py

@@ -406,17 +406,69 @@ def test_compute_reward_loss_identical_sequences():
 
     chosen_batch = SimpleNamespace(
         input_ids=input_ids,
-        rewards=torch.randn(1, seq_len),
+        rewards=torch.randn(1, seq_len, requires_grad=True),
     )
     rejected_batch = SimpleNamespace(
         input_ids=input_ids.clone(),
-        rewards=torch.randn(1, seq_len),
+        rewards=torch.randn(1, seq_len, requires_grad=True),
     )
     loss = GPT2RewardModel.compute_reward_loss(
         chosen_batch, rejected_batch, pad_token_id=pad_token_id
     )
     assert loss.shape == torch.Size([])
     assert loss.item() == 0.0
+    loss.backward()
+    torch.testing.assert_close(
+        chosen_batch.rewards.grad, torch.zeros_like(chosen_batch.rewards)
+    )
+    torch.testing.assert_close(
+        rejected_batch.rewards.grad, torch.zeros_like(rejected_batch.rewards)
+    )
+
+
+def test_compute_reward_loss_normalizes_by_non_identical_sequences():
+    pad_token_id = 50256
+    chosen_ids = torch.tensor(
+        [
+            [1, 2, 3, 4, pad_token_id],
+            [1, 2, 9, 4, pad_token_id],
+        ]
+    )
+    rejected_ids = torch.tensor(
+        [
+            [1, 2, 3, 4, pad_token_id],
+            [1, 2, 3, 4, pad_token_id],
+        ]
+    )
+    chosen_rewards = torch.tensor(
+        [
+            [0.0, 0.0, 10.0, 10.0, 0.0],
+            [0.0, 0.0, 2.0, 2.0, 0.0],
+        ],
+        requires_grad=True,
+    )
+    rejected_rewards = torch.tensor(
+        [
+            [0.0, 0.0, -10.0, -10.0, 0.0],
+            [0.0, 0.0, 1.0, 1.0, 0.0],
+        ],
+        requires_grad=True,
+    )
+    chosen_batch = SimpleNamespace(input_ids=chosen_ids, rewards=chosen_rewards)
+    rejected_batch = SimpleNamespace(input_ids=rejected_ids, rewards=rejected_rewards)
+
+    loss = GPT2RewardModel.compute_reward_loss(
+        chosen_batch, rejected_batch, pad_token_id=pad_token_id
+    )
+    expected_loss = -F.logsigmoid(chosen_rewards[1, 2:4] - rejected_rewards[1, 2:4])
+    torch.testing.assert_close(loss, expected_loss.mean())
+    loss.backward()
+    torch.testing.assert_close(
+        chosen_rewards.grad[0], torch.zeros_like(chosen_rewards[0])
+    )
+    torch.testing.assert_close(
+        rejected_rewards.grad[0], torch.zeros_like(rejected_rewards[0])
+    )
 
 
 @pytest.mark.skipif(

torchrl/modules/models/llm.py

@@ -122,7 +122,8 @@ def compute_reward_loss(chosen_batch, rejected_batch, pad_token_id=50256):
         rejected_rewards = rejected_batch.rewards
 
         bs = chosen_rewards.shape[0]
-        loss = torch.tensor(0.0, device=chosen_rewards.device)
+        loss = None
+        valid_count = 0
 
         # TODO: this loop can likely be made more efficient
         for i in range(bs):
@@ -144,8 +145,12 @@ def compute_reward_loss(chosen_batch, rejected_batch, pad_token_id=50256):
             c_truncated_reward = chosen_rewards[i][divergence_ind:end_ind]
             r_truncated_reward = rejected_rewards[i][divergence_ind:end_ind]
 
-            loss += -F.logsigmoid(c_truncated_reward - r_truncated_reward).mean()
-        return loss / bs
+            sample_loss = -F.logsigmoid(c_truncated_reward - r_truncated_reward).mean()
+            loss = sample_loss if loss is None else loss + sample_loss
+            valid_count += 1
+        if loss is None:
+            return chosen_rewards.sum() * 0.0 + rejected_rewards.sum() * 0.0
+        return loss / valid_count
 
     @classmethod
     def from_pretrained(cls, path):