Repository: huggingface/trl

Issue:
DPOTrainer ignores training & evaluation data if there are more than 2 responses
Hello!

First of all, thank you @younesbelkada and @lvwerra for maintaining this work, and thank you @kashif for spearheading the valuable DPO implementation.

## Bug overview
* The `DPODataCollatorWithPadding` processes `batch_size` samples at once, yielding `batch_size` chosen-rejected pairs. However, `batch_size` samples may have much more than `batch_size` pairs.

## Details
The issue originates in the `DPODataCollatorWithPadding.__call__` method:
https://github.com/lvwerra/trl/blob/7f297b38c665e6d2357dfd5ba552ef685afe1302/trl/trainer/utils.py#L328-L345
The `break` will pre-emptively stop processing of all pairs in `features` to ensure that the output has a length of `batch_size`. Because of this, any training or evaluation data with more than 2 responses will result in data not being used.

## How to reproduce
<details><summary>See reproduction steps</summary>

Replace `DPODataCollatorWithPadding.__call__` with this variant, which tracks how many samples are used and how many are ignored:
```python
    def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
        batch = []
        done = False
        for feature in features:
            prompt = feature["prompt"]
            responses = feature["responses"]
            pairs = feature["pairs"]

            for p in pairs:
                if not done:

                    batch_element = self.tokenize_batch_element(prompt, responses[p[0]], responses[p[1]])
                    batch.append(batch_element)

                    done = len(batch) == self.batch_size
                    self.included += 1
                else:
                    self.ignored += 1

        print(f"{self.ignored / (self.ignored + self.included):%} of training pairs are ignored.")
        return self.collate(batch)
```
and add
```python
    included: int = 0
    ignored: int = 0
```
to the list of dataclass fields of the `DPODataCollatorWithPadding`.

Then, run this script, which uses a dummy dataset with many responses:
```python
from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments
from trl import DPOTrainer
from datasets import Dataset

model = AutoModelForCausalLM.from_pretrained("gpt2")

model_ref = AutoModelForCausalLM.from_pretrained("gpt2")
tokenizer = AutoTokenizer.from_pretrained("gpt2")
tokenizer.pad_token = tokenizer.eos_token

train_dataset_dict = {
    "prompt": [
        "trl",
        "is",
        "pretty",
        "great",
        "but",
        "it",
        "has",
        "a",
        "small",
        "bug",
    ],
    "responses": [
        [
            "good",
            "bad",
            "worse",
            "horrible",
            "not good",
            "poor",
            "awful",
            "substandard",
            "unsatisfactory",
        ],
    ]
    * 10,
    "pairs": [
        [(0, 1), (0, 2), (0, 3), (0, 4), (0, 5), (0, 6), (0, 7), (0, 8)],
    ]
    * 10,
}
train_dataset = Dataset.from_dict(train_dataset_dict)

training_args = TrainingArguments(
    per_device_train_batch_size=4,
    remove_unused_columns=False,
    output_dir="./tmp",
    report_to="none",
    num_train_epochs=10,
)

dpo_trainer = DPOTrainer(
    model,
    model_ref,
    args=training_args,
    beta=0.1,
    train_dataset=train_dataset,
    tokenizer=tokenizer,
)

dpo_trainer.train()
```
Towards the end of training, you'll see that `85.00%` of all pairs have been ignored in this example. 

</details>

## Proposed fixes
Before I start working on it, I'd like to discuss my proposed changes: 

### The fix
Spread out the training and evaluation datasets into a Dataset of `prompt`, `chosen`, `rejected` columns right in the `DPOTrainer.__init__`. This requires some small changes throughout this work, but will work conveniently otherwise: each dataset element will be one training sample. It also has a drawback: prompts and responses will need to be tokenized multiple times, unless we immediately store the `prompt`-`chosen`-`rejected` dataset in a tokenized way.

### Alternative input structures
#### `prompt`, `chosen`, `rejected`
Beyond that, I can also implement a more drastic change: require the user to provide a Dataset with `prompt`, `chosen`, `rejected` columns. 

#### `prompt`, `responses`, `chosen` | `pairs`
A less drastic change is to allow users to provide a `chosen` column like:
```
"chosen": [0, 1, 0, 0]
```
where the `chosen` is simply the index of the best response, instead of `pairs` like:
```python
"pairs": [
        [(0, 1)],
        [(1, 0)],
        [(0, 2), (0, 1)],
        [(0, 1), (0, 2), (0, 3)],
    ]
```
Internally, we can use this `chosen` to create the same pairs. This is simpler for most users, but is less powerful, as the `pairs` approach allows for a ranking (e.g. also say that although response 1 is best, response 2 is better than response 3). I propose that the input dataset should *either* have a `pairs` or `chosen` column.

### What now?
I would love some feedback on your thoughts on:
* Is the proposed fix (internal `prompt`, `chosen`, `rejected` columns) a good idea?
* Should we change the input structures to `prompt`, `chosen`, `rejected`, `prompt`, `responses`, `chosen` | `pairs`, or keep it as-is?

---

- Tom Aarsen


Discussion:
thanks @tomaarsen for the report. I know of the batch issue but didn't realize it would be so bad for this use case... let me review your kind proposal and brainstorm. in the original code, the batch collation occurs on the dataset preparation side...
Hi @tomaarsen! I'm new to the repository and I'm wondering if you could help me understand the issue you identified. I thought the current datasets only have 2 responses per prompt? (the chosen response and the rejected response). That is, I thought len(pairs) = len(responses) = 2 for all batches. In that case, the code should work fine, or did I misunderstand anything?

Thanks!
@liutianlin0121 no the dataset can potentially have more than one preference per prompt and in the collator, i am looping over each prompt and when the result has the required batch size i return from that prompt/preference example and move onto the next, meaning that for an example where there are many pairs (say more than the training batch size), currently the collator is just not using them... 
> Hi @tomaarsen! I'm new to the repository and I'm wondering if you could help me understand the issue you identified. I thought the current datasets only have 2 responses per prompt? (the chosen response and the rejected response). That is, I thought len(pairs) = len(responses) = 2 for all batches. In that case, the code should work fine, or did I misunderstand anything?
> 
> Thanks!

@kashif's response is correct. Here's a bit more information that should clear it right up: https://huggingface.co/docs/trl/main/en/dpo_trainer#expected-dataset-format

- Tom Aarsen
@tomaarsen  @liutianlin0121  at the moment I am tending towards setting the default data format for the DPO to be "prompt", "chosen", and "rejected" keys and let the individual dataset helpers create these dicts. Then I believe everything should work with the DPO collator... let me send you a PR for review?  It is your first suggestion but I believe the change is only in the docs and in the individual dataset helper, and not in the DPOTrainer... let me check 
Thanks both! 👍

Resolving patch:
diff --git a/examples/dpo.py b/examples/dpo.py
--- a/examples/dpo.py
+++ b/examples/dpo.py
@@ -1,10 +1,8 @@
 # 0. imports
-from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import Optional
+from typing import Dict, Optional
 
 import torch
-import tqdm
 from datasets import Dataset, load_dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer, HfArgumentParser, TrainingArguments
 
@@ -66,8 +64,8 @@ def get_hh(split: str, sanity_check: bool = False, silent: bool = False, cache_d
     The dataset is converted to a dictionary with the following structure:
     {
         'prompt': List[str],
-        'responses': List[List[str]],
-        'pairs': List[Tuple[int, int]]
+        'chosen': List[str],
+        'rejected': List[str],
     }
 
     Prompts should be structured as follows:
@@ -78,30 +76,15 @@ def get_hh(split: str, sanity_check: bool = False, silent: bool = False, cache_d
     if sanity_check:
         dataset = dataset.select(range(min(len(dataset), 1000)))
 
-    def split_prompt_and_responses(ex):
-        prompt = extract_anthropic_prompt(ex["chosen"])
-        chosen_response = ex["chosen"][len(prompt) :]
-        rejected_response = ex["rejected"][len(prompt) :]
-        return prompt, chosen_response, rejected_response
-
-    data = defaultdict(lambda: defaultdict(list))
-    for row in tqdm.tqdm(dataset, desc="Processing HH", disable=silent):
-        prompt, chosen, rejected = split_prompt_and_responses(row)
-        responses = [chosen, rejected]
-        n_responses = len(data[prompt]["responses"])
-        data[prompt]["pairs"].append((n_responses, n_responses + 1))
-        data[prompt]["responses"].extend(responses)
-        data[prompt]["sft_target"] = chosen
-
-    def gen():
-        for prompt, values in data.items():
-            yield {
-                "prompt": prompt,
-                "responses": values["responses"],
-                "pairs": values["pairs"],
-            }
-
-    return Dataset.from_generator(gen)
+    def split_prompt_and_responses(sample) -> Dict[str, str]:
+        prompt = extract_anthropic_prompt(sample["chosen"])
+        return {
+            "prompt": prompt,
+            "chosen": sample["chosen"][len(prompt) :],
+            "rejected": sample["rejected"][len(prompt) :],
+        }
+
+    return dataset.map(split_prompt_and_responses)
 
 
 if __name__ == "__main__":
diff --git a/trl/trainer/utils.py b/trl/trainer/utils.py
--- a/trl/trainer/utils.py
+++ b/trl/trainer/utils.py
@@ -326,23 +326,18 @@ def collate(self, batch):
         return padded_batch
 
     def __call__(self, features: List[Dict[str, Any]]) -> Dict[str, Any]:
-        batch = []
+        tokenized_batch = []
+
         for feature in features:
             prompt = feature["prompt"]
-            responses = feature["responses"]
-            pairs = feature["pairs"]
-
-            done = False
-            for p in pairs:
-                if done:
-                    break
-
-                batch_element = self.tokenize_batch_element(prompt, responses[p[0]], responses[p[1]])
-                batch.append(batch_element)
+            chosen = feature["chosen"]
+            rejected = feature["rejected"]
 
-                done = len(batch) == self.batch_size
+            batch_element = self.tokenize_batch_element(prompt, chosen, rejected)
+            tokenized_batch.append(batch_element)
 
-        return self.collate(batch)
+        # return collated batch
+        return self.collate(tokenized_batch)
 
 
 class ConstantLengthDataset(IterableDataset):


Test patch:
diff --git a/tests/test_dpo_trainer.py b/tests/test_dpo_trainer.py
--- a/tests/test_dpo_trainer.py
+++ b/tests/test_dpo_trainer.py
@@ -48,19 +48,28 @@ def test_dpo_trainer(self):
                     "hello",
                     "how are you",
                     "What is your name?",
+                    "What is your name?",
+                    "Which is the best programming language?",
+                    "Which is the best programming language?",
                     "Which is the best programming language?",
                 ],
-                "responses": [
-                    ["hi nice to meet you", "leave me alone"],
-                    ["I am not fine", "I am fine"],
-                    ["My name is Mary", "Whats it to you?", "I dont have a name"],
-                    ["Python", "Javascript", "C++", "Java"],
+                "chosen": [
+                    "hi nice to meet you",
+                    "I am fine",
+                    "My name is Mary",
+                    "My name is Mary",
+                    "Python",
+                    "Python",
+                    "Python",
                 ],
-                "pairs": [
-                    [(0, 1)],
-                    [(1, 0)],
-                    [(0, 2), (0, 1)],
-                    [(0, 1), (0, 2), (0, 3)],
+                "rejected": [
+                    "leave me alone",
+                    "I am not fine",
+                    "Whats it to you?",
+                    "I dont have a name",
+                    "Javascript",
+                    "C++",
+                    "Java",
                 ],
             }
             # fmt: on