Repository: huggingface/trl
Pull request: Relax reward trainer constraint
Fixes https://github.com/lvwerra/trl/issues/535

We should relax the constraint for users in case they implement their own Data collator when using `RewardTrainer`.

cc @lvwerra 

Problem statement:
Use my own compute_loss with RewardTrainer
Hi. I want to use my own RewardModel (not inherited from AutoModelForSequenceClassification)

But, i got message like below.
```
compute_loss is only implemented for RewardDataCollatorWithPadding, please implement your own compute_loss method if you are using a custom data collator
```

Is there any parameter which i can pass my own compute_loss to RewardTrainer ?

Hints from issue discussion:
Hi @allzero-kwon, it looks like you passed a custom data collator which is why you get this message. @younesbelkada maybe we don't need to through an error here, no? since a user might pass a viable alternative to the default collator that does not require writing a custom loss?

In any case, if you want to write your own loss function you can do it as follows:
```python
from trl import RewardTrainer
class MyRewardTrainer(RewardTrainer):
    def compute_loss(self, model, inputs):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        logits = outputs[0]
        return my_custom_loss(logits, labels)
``` 

Hope this helps!

Edit: here the link to the line causing the issue:
https://github.com/lvwerra/trl/blob/7f297b38c665e6d2357dfd5ba552ef685afe1302/trl/trainer/reward_trainer.py#L163

Commit message: relax reward trainer constraint

Patch to review:
diff --git a/trl/trainer/reward_trainer.py b/trl/trainer/reward_trainer.py
index 8649dfb502..44c963d52c 100644
--- a/trl/trainer/reward_trainer.py
+++ b/trl/trainer/reward_trainer.py
@@ -161,8 +161,10 @@ def compute_loss(
         return_outputs=False,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
         if not self.use_reward_data_collator:
-            raise NotImplementedError(
-                "compute_loss is only implemented for RewardDataCollatorWithPadding, please implement your own compute_loss method if you are using a custom data collator"
+            warnings.warn(
+                "The current compute_loss is implemented for RewardDataCollatorWithPadding"
+                " if you are using a custom data collator make sure you know what you are doing or"
+                " implement your own compute_loss method."
             )
         rewards_chosen = model(
             input_ids=inputs["input_ids_chosen"],


Code review comments:
diff_hunk: @@ -161,8 +161,10 @@ def compute_loss(
         return_outputs=False,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, Dict[str, torch.Tensor]]]:
         if not self.use_reward_data_collator:
-            raise NotImplementedError(
-                "compute_loss is only implemented for RewardDataCollatorWithPadding, please implement your own compute_loss method if you are using a custom data collator"
+            warnings.warn(
+                "The current compute_loss is implemented for RewardDataCollatorWithPadding"
path: trl/trainer/reward_trainer.py
text: @user1:
```suggestion
                "The current compute_loss is implemented for RewardDataCollatorWithPadding,"
```