Repository: huggingface/trl

Issue:
Potential bug of model output dim in reward trainer example
Hi!

I've been exploring the `examples/scripts/reward_trainer.py` and believe I've found a potential bug that I'd like to address.

## Potential bug.
To my understanding, the reward model should produce a scalar output. However, when we load the pre-trained model in the [example script](https://github.com/lvwerra/trl/blob/main/examples/scripts/reward_trainer.py#L75-L80), the model will be a binary classifier by default, with a 2-dim logit output:

```python
import torch
inputs = tokenizer("Hello, my dog is cute", return_tensors="pt")
model = model.cpu()
with torch.no_grad():
    out = model(**inputs)
print(out.logits.shape) # torch.Size([1, 2])
```

During training, this issue was hidden by the [the loss computation part](https://github.com/lvwerra/trl/blob/main/trl/trainer/reward_trainer.py#L169-L176) of the reward trainer:

```python
    rewards_chosen = model(
        input_ids=inputs["input_ids_chosen"],
        attention_mask=inputs["attention_mask_chosen"],
    )[0]

    rewards_rejected = model(
        input_ids=inputs["input_ids_rejected"],
        attention_mask=inputs["attention_mask_rejected"],
    )[0]

    loss = -torch.nn.functional.logsigmoid(rewards_chosen - rewards_rejected).mean()
```
There, the `rewards_rejected` and `rewards_chosen` are dim [sample, 2] matrices, which I think should be dim [sample, 1] instead. The `.mean()` operator hides the issue as it returns a scalar value.

Here is a [colab](https://colab.research.google.com/drive/1szzGC8H1DQ7329_idVnnOnPIPE6D7J4j?usp=sharing) that shows the dimensionality problem.

## Proposed fix

To resolve this, in the model training example, we can use:

```python
model = AutoModelForSequenceClassification.from_pretrained(
    script_args.model_name,
    quantization_config=quantization_config,
    device_map=device_map,
    trust_remote_code=script_args.trust_remote_code,
    num_labels=1 # <---- make sure that the model output is a scalar
)
```
I'll submit a pull request if it is OK. 

In addition, while exploring `examples/scripts/reward_trainer.py`, I noticed a few other places that can be improved. For example, there is currently no validation steps that show the genearlization performance of the reward model; the `log_with` option in `ScriptArguments` is currently not being used anywhere and therefore has no effect. I'm happy to include the validation and fix the logging glitch in my pull request as well.

Tianlin


Discussion:
Hi @liutianlin0121 
Thanks for the issue, I went through the shared materials and details. I also have checked with @edbeeching and we think that you are right, the proposed fix seems to be the correct one. Would you ne happy opening a PR for the fix?
Thanks! Yes I'll submit a PR soon :-) 

Resolving patch:
diff --git a/examples/scripts/reward_trainer.py b/examples/scripts/reward_trainer.py
--- a/examples/scripts/reward_trainer.py
+++ b/examples/scripts/reward_trainer.py
@@ -43,8 +43,13 @@ class ScriptArguments:
     dataset_name: Optional[str] = field(default="Anthropic/hh-rlhf", metadata={"help": "the model name"})
     dataset_text_field: Optional[str] = field(default="text", metadata={"help": "the text field of the dataset"})
     log_with: Optional[str] = field(default=None, metadata={"help": "use 'wandb' to log with wandb"})
+    logging_steps: Optional[int] = field(default=500, metadata={"help": "the number of update steps between two logs"})
+    eval_split: Optional[str] = field(
+        default="none", metadata={"help": "the dataset split to evaluate on; default to 'none' (no evaluation)"}
+    )
     learning_rate: Optional[float] = field(default=1.41e-5, metadata={"help": "the learning rate"})
     batch_size: Optional[int] = field(default=64, metadata={"help": "the batch size"})
+    num_train_epochs: Optional[int] = field(default=1, metadata={"help": "the number of training epochs"})
     seq_length: Optional[int] = field(default=512, metadata={"help": "Input sequence length"})
     gradient_accumulation_steps: Optional[int] = field(
         default=16, metadata={"help": "the number of gradient accumulation steps"}
@@ -77,11 +82,14 @@ class ScriptArguments:
     quantization_config=quantization_config,
     device_map=device_map,
     trust_remote_code=script_args.trust_remote_code,
+    num_labels=1,
 )
 
 # Step 2: Load the dataset and pre-process it
 tokenizer = AutoTokenizer.from_pretrained(script_args.model_name)
-dataset = load_dataset(script_args.dataset_name, split="train")
+train_dataset = load_dataset(script_args.dataset_name, split="train")
+
+
 # Turn the dataset into pairs of post + summaries, where text_j is the preferred question + answer and text_k is the other.
 # Then tokenize the dataset.
 # Adapt this section to your needs for custom datasets
@@ -107,8 +115,7 @@ def preprocess_function(examples):
 
 
 # preprocess the dataset and filter out QAs that are longer than script_args.max_length
-original_columns = dataset.column_names
-train_dataset = dataset.map(
+train_dataset = train_dataset.map(
     preprocess_function,
     batched=True,
     num_proc=4,
@@ -118,13 +125,34 @@ def preprocess_function(examples):
     and len(x["input_ids_rejected"]) <= script_args.seq_length
 )
 
+if script_args.eval_split == "none":
+    eval_dataset = None
+else:
+    eval_dataset = load_dataset(script_args.dataset_name, split=script_args.eval_split)
+
+    eval_dataset = eval_dataset.map(
+        preprocess_function,
+        batched=True,
+        num_proc=4,
+    )
+    eval_dataset = eval_dataset.filter(
+        lambda x: len(x["input_ids_chosen"]) <= script_args.seq_length
+        and len(x["input_ids_rejected"]) <= script_args.seq_length
+    )
+
 
 # Step 3: Define the training arguments
 training_args = TrainingArguments(
     output_dir=script_args.output_dir,
     per_device_train_batch_size=script_args.batch_size,
+    num_train_epochs=script_args.num_train_epochs,
     gradient_accumulation_steps=script_args.gradient_accumulation_steps,
     learning_rate=script_args.learning_rate,
+    report_to="wandb" if script_args.log_with == "wandb" else "tensorboard",
+    remove_unused_columns=False,
+    optim="adamw_torch",
+    logging_steps=script_args.logging_steps,
+    evaluation_strategy="steps" if script_args.eval_split != "none" else "no",
 )
 
 # Step 4: Define the LoraConfig
@@ -139,6 +167,7 @@ def preprocess_function(examples):
     tokenizer=tokenizer,
     args=training_args,
     train_dataset=train_dataset,
+    eval_dataset=eval_dataset,
     peft_config=peft_config,
     max_length=script_args.seq_length,
 )


Test patch: