Repository: huggingface/trl
Pull request: Update incorrect data processing in DataCollatorForChatML
# What does this PR do?
Fix the extra BOS token and the absence of an EOS token in the returned input_ids, and potentially the absence of a target string in the returned labels.

<!--
Congratulations! You've made it this far! You're not quite done yet though.

Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.

Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.

Once you're done, someone will review your PR shortly. They may suggest changes to make the code even better.
-->

<!-- Remove if not applicable -->

Fixes #2169


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/trl/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [x] Was this discussed/approved via a GitHub issue? Please add a link
      to it if that's the case. (https://github.com/huggingface/trl/issues/2169)
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/trl/tree/main/docs).
- [x] Did you write any new necessary tests?


## Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

Problem statement:
Incorrect data processing in DataCollatorForChatML
### System Info

Python 3.11.9
trl 0.11.0
transformers 4.45.1

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder
- [X] My own task or dataset (give details below)

### Reproduction

```python
from trl.trainer.utils import DataCollatorForChatML
from transformers import AutoTokenizer
# base model settings
tokenizer = AutoTokenizer.from_pretrained("codellama/CodeLlama-7b-Instruct-hf")
tokenizer.pad_token = tokenizer.bos_token if tokenizer.pad_token is None else tokenizer.pad_token

# Initialize the data collator for processing the chat model inputs
max_length = 1024
ignore_index = -100
messages_key = "messages"
examples = [
    {
        messages_key:
            [
                {"role": "user", "content": "Does the following code contain any security vulnerabilities? Return true or false.\nchar buffer[10];\nchar input[50];\nstrcpy(buffer, input);\n"}, 
                {"role": "assistant", "content": "true"}
            ]
    }
]
collator = DataCollatorForChatML(tokenizer=tokenizer, max_length=max_length, ignore_index=ignore_index, messages_key=messages_key)

# Process the data
data = collator(examples)
print(data["input_ids"])
print(data["labels"])
```

# Founded Bugs and Solutions

## Output from Demo Script

```python
>>> print(data["input_ids"])
tensor([[    1,     1,   518, 25580, 29962,  5538,   278,  1494,   775,  1712,
           738,  6993, 23180, 11614, 29973,  7106,  1565,   470,  2089, 29889,
            13,  3090,  6835, 29961, 29896, 29900,  1385,    13,  3090,  1881,
         29961, 29945, 29900,  1385,    13,   710, 23141, 29898,  9040, 29892,
          1881,   416,   518, 29914, 25580, 29962, 29871,     1,  1565]])
>>> print(data["labels"])
tensor([[-100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,
         -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,
         -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,
         -100, -100, -100, -100, -100, -100, -100, -100, -100, -100, -100,    1,
            2]])
```

## Founded Bugs

Preliminary: the expected output from assistant is a string "true" and its token id is 1565. The bos token id is 1 and the eos token id is 2.

1. Extra bos token and no eos token in data["input_ids"]. We can see that data["input_ids"] got an extra bos token at the beginning (i.e., [1, 1, ...]), and at the end of data["input_ids"] we can see another extra bos token (i.e., [..., 1, 1565]) and no eos token.
2. No target string in data["labels"]. As for the data["labels"], it does not include the expected output from assistant ("true" token) and only the bos and eos token is preserved (i.e., [..., -100, 1, 2]).

## Expected Output of Demo Script

The expected output should look like this:

```python
>>> print(data["input_ids"])
tensor([[    1,   518, 25580, 29962,  5538,   278,  1494,   775,  1712,   738,
          6993, 23180, 11614, 29973,  7106,  1565,   470,  2089, 29889,    13,
          3090,  6835, 29961, 29896, 29900,  1385,    13,  3090,  1881, 29961,
         29945, 29900,  1385,    13,   710, 23141, 29898,  9040, 29892,  1881,
           416,   518, 29914, 25580, 29962, 29871,  1565, 29871,     2]])
>>> print(data["labels"])
tensor([[ -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
          -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,  -100,
          -100,  -100,  -100,  -100,  -100,  -100,  1565, 29871,     2]])
```

We can see the expected output of data["input_ids"] and data["labels"] should not have any extra bos tokens, and the data["labels"] preserved the target string correctly (i.e., [..., -100, 1565, 29871, 2]).
I believe this result is more compatible with the loss calculation like LlamaForCausalLM as demonstrated below:

```python
loss = None
if labels is not None:
    # Upcast to float if we need to compute the loss to avoid potential precision issues
    logits = logits.float()
    # Shift so that tokens < n predict n
    shift_logits = logits[..., :-1, :].contiguous()
    shift_labels = labels[..., 1:].contiguous()
    # Flatten the tokens
    loss_fct = CrossEntropyLoss()
    shift_logits = shift_logits.view(-1, self.config.vocab_size)
    shift_labels = shift_labels.view(-1)
    # Enable model parallelism
    shift_labels = shift_labels.to(shift_logits.device)
    loss = loss_fct(shift_logits, shift_labels)
```

## Location and Solution to the Bugs

The bugs is located at the beginning of __call__ function of DataCollatorForChatML:

```python
prompts = []
completions = []

for example in examples:
    messages = example[self.messages_key]
    formatted_chat = self.tokenizer.apply_chat_template(messages, tokenize=False)

    # Split the formatted chat into prompt and completion
    assistant_messages = [msg for msg in messages if msg["role"] == "assistant"]
    last_assistant_message = assistant_messages[-1]["content"]
    prompt = formatted_chat.rsplit(last_assistant_message, 1)[0]
    completion = last_assistant_message

    prompts.append(prompt)
    completions.append(completion)

# Tokenize prompts and completions
tokenized_prompts = self.tokenizer(
    prompts, truncation=True, max_length=self.max_length, padding=False, return_tensors=None
)
tokenized_completions = self.tokenizer(
    completions, truncation=True, max_length=self.max_length, padding=False, return_tensors=None
)

# Combine prompts and completions
input_ids = []
attention_mask = []
labels = []

for prompt, completion in zip(tokenized_prompts["input_ids"], tokenized_completions["input_ids"]):
    combined_input_ids = prompt + completion
    combined_attention_mask = [1] * len(combined_input_ids)

    # Create labels for one-token ahead task, masking the prompt
    combined_labels = [self.ignore_index] * len(prompt) + completion[:-1]
    combined_labels.append(self.tokenizer.eos_token_id)  # Add EOS token as final target

    input_ids.append(combined_input_ids)
    attention_mask.append(combined_attention_mask)
    labels.append(combined_labels)
```

### Reason for Bug1

The code that raises bug1 is due to:

```python
# Tokenize prompts and completions
tokenized_prompts = self.tokenizer(
    prompts, truncation=True, max_length=self.max_length, padding=False, return_tensors=None
)
tokenized_completions = self.tokenizer(
    completions, truncation=True, max_length=self.max_length, padding=False, return_tensors=None
)
```

Because prompts are obtained from formatted_chat, and formatted_chat already adds bos token by default:

```python
formatted_chat = self.tokenizer.apply_chat_template(messages, tokenize=False)
```

So, after using self.tokenizer to tokenize prompts and completions, it will add extra bos token to them.

### Fix to Bug1

Set the default parameter of add_special_token from True to False in the following code:

```python
# Tokenize prompts and completions
tokenized_prompts = self.tokenizer(
    prompts, truncation=True, max_length=self.max_length, padding=False, return_tensors=None, add_special_token=False
)
tokenized_completions = self.tokenizer(
    completions, truncation=True, max_length=self.max_length, padding=False, return_tensors=None, add_special_token=False
)
```

### Reason for Bug2

The location that raises bug2 is here:

```python
completion = last_assistant_message
```

Since apply_chat_template will add eos token to formatted_chat, above code does not include the eos token to the completion correctly. And 

```python
tokenized_completions = self.tokenizer(
    completions, truncation=True, max_length=self.max_length, padding=False, return_tensors=None
)
```

will adds extra bos token to the tokenized_completions, so in this code:

```python
combined_labels = [self.ignore_index] * len(prompt) + completion[:-1]
```

the expected output from assistant is ignored and only the bos token is retained.

### Fix to Bug2

The correct implementation should be:

```python
completion = last_assistant_message + formatted_chat.rsplit(last_assistant_message, 1)[1].
```

Hints from issue discussion:
Thanks a lot for this detailed report @ruijunfeng 
This is indeed a critical issue. Are you willing to submit a PR to solve it?
> Thanks a lot for this detailed report @ruijunfeng This is indeed a critical issue. Are you willing to submit a PR to solve it?

Hi there, I have submitted a PR to fix this, hope this will help😊

Commit message: Update tests/test_utils.py

Patch to review:
diff --git a/tests/test_utils.py b/tests/test_utils.py
index e79edc755f..dbd68e0b44 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -20,7 +20,13 @@
 from transformers.utils import is_peft_available
 
 from trl.trainer.model_config import ModelConfig
-from trl.trainer.utils import decode_and_strip_padding, generate_model_card, get_peft_config, pad
+from trl.trainer.utils import (
+    DataCollatorForChatML,
+    decode_and_strip_padding,
+    generate_model_card,
+    get_peft_config,
+    pad,
+)
 
 
 if is_peft_available():
@@ -169,3 +175,87 @@ def test_val_none(self):
         assert "my_model" in card_text
         assert 'pipeline("text-generation", model="username/my_hub_model", device="cuda")' in card_text
         assert "My Trainer" in card_text
+
+
+class TestDataCollatorForChatML(unittest.TestCase):
+    def setUp(self):
+        # Initialize the tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained("codellama/CodeLlama-7b-Instruct-hf")
+        self.tokenizer.pad_token = (
+            self.tokenizer.bos_token if self.tokenizer.pad_token is None else self.tokenizer.pad_token
+        )
+
+        # Define token IDs
+        self.bos_token_id = self.tokenizer.bos_token_id if self.tokenizer.bos_token_id is not None else 1
+        self.eos_token_id = self.tokenizer.eos_token_id if self.tokenizer.eos_token_id is not None else 2
+        self.assistant_output_token_id = 1565  # Token ID for "true", which is the last assistant's response in the example
+        self.ignore_index = -100
+        self.max_length = 1024
+        self.messages_key = "messages"
+
+        # Example input
+        self.examples = [
+            {
+                self.messages_key: [
+                    {
+                        "role": "user",
+                        "content": (
+                            "Does the following code contain any security vulnerabilities? Return true or false.\n"
+                            "char buffer[10];\nchar input[50];\nstrcpy(buffer, input);\n"
+                        ),
+                    },
+                    {"role": "assistant", "content": "true"},
+                ]
+            }
+        ]
+
+        # Initialize the data collator
+        self.collator = DataCollatorForChatML(
+            tokenizer=self.tokenizer,
+            max_length=self.max_length,
+            ignore_index=self.ignore_index,
+            messages_key=self.messages_key,
+        )
+
+    def test_data_collator_for_chatml(self):
+        # Process the data
+        data = self.collator(self.examples)
+
+        # Decode input_ids and labels for verification
+        input_ids = data["input_ids"][0].tolist()
+        labels = data["labels"][0].tolist()
+
+        # Expected tokens
+        expected_bos = self.bos_token_id
+        expected_eos = self.eos_token_id
+        expected_assistant_token = self.assistant_output_token_id
+
+        # Verify that input_ids start with a BOS token and there are no extra ones
+        self.assertEqual(input_ids[0], expected_bos, "The first token of input_ids should be BOS token.")
+        self.assertNotEqual(
+            input_ids[1], expected_bos, "The second token of input_ids should not be BOS token (extra BOS)."
+        )
+
+        # Verify that the assistant's response token is present in input_ids
+        self.assertIn(expected_assistant_token, input_ids, "Assistant's response token should be in input_ids.")
+
+        # Verify that EOS token is at the end of input_ids
+        self.assertEqual(input_ids[-1], expected_eos, "The last token of input_ids should be EOS token.")
+
+        # Verify that the labels preserved the target string (last_assistant_response)
+        last_assistant_response = self.examples[0][self.messages_key][-1]["content"]
+        last_assistant_response_tokens = self.tokenizer.encode(last_assistant_response, add_special_tokens=False)
+
+        # Find the start and end of the last assistant's response in the labels
+        response_start = next(i for i, label in enumerate(labels) if label != self.ignore_index)
+        response_end = next(i for i in range(len(labels) - 1, -1, -1) if labels[i] != self.ignore_index)
+
+        actual_response = labels[response_start : response_end - 1]
+        self.assertEqual(
+            actual_response,
+            last_assistant_response_tokens,
+            "The labels should preserve the last assistant's response tokens.",
+        )
+
+        # Verify that EOS token is at the end of labels
+        self.assertEqual(labels[-1], expected_eos, "The last token of labels should be EOS token.")
diff --git a/trl/trainer/utils.py b/trl/trainer/utils.py
index 1591b75ada..f09fab23a2 100644
--- a/trl/trainer/utils.py
+++ b/trl/trainer/utils.py
@@ -265,17 +265,29 @@ def __call__(self, examples: List[Dict[str, Any]]) -> Dict[str, torch.Tensor]:
             assistant_messages = [msg for msg in messages if msg["role"] == "assistant"]
             last_assistant_message = assistant_messages[-1]["content"]
             prompt = formatted_chat.rsplit(last_assistant_message, 1)[0]
-            completion = last_assistant_message
+            completion = last_assistant_message + formatted_chat.rsplit(last_assistant_message, 1)[1]
 
             prompts.append(prompt)
             completions.append(completion)
 
         # Tokenize prompts and completions
         tokenized_prompts = self.tokenizer(
-            prompts, truncation=True, max_length=self.max_length, padding=False, return_tensors=None
+            prompts,
+            truncation=True,
+            max_length=self.max_length,
+            padding=False,
+            return_tensors=None,
+            # We assume the inputs are already wrapped with BOS&EOS tokens in tokenizer.apply_chat_template, so extra BOS/EOS tokens should not be added
+            add_special_tokens=False,
         )
         tokenized_completions = self.tokenizer(
-            completions, truncation=True, max_length=self.max_length, padding=False, return_tensors=None
+            completions,
+            truncation=True,
+            max_length=self.max_length,
+            padding=False,
+            return_tensors=None,
+            # We assume the inputs are already wrapped with BOS&EOS tokens in tokenizer.apply_chat_template, so extra BOS/EOS tokens should not be added
+            add_special_tokens=False,
         )
 
         # Combine prompts and completions


Code review comments:
diff_hunk: @@ -169,3 +175,87 @@ def test_val_none(self):
         assert "my_model" in card_text
         assert 'pipeline("text-generation", model="username/my_hub_model", device="cuda")' in card_text
         assert "My Trainer" in card_text
+
+
+class TestDataCollatorForChatML(unittest.TestCase):
+    def setUp(self):
+        # Initialize the tokenizer
+        self.tokenizer = AutoTokenizer.from_pretrained("codellama/CodeLlama-7b-Instruct-hf")
+        self.tokenizer.pad_token = (
+            self.tokenizer.bos_token if self.tokenizer.pad_token is None else self.tokenizer.pad_token
+        )
+
+        # Define token IDs
+        self.bos_token_id = self.tokenizer.bos_token_id if self.tokenizer.bos_token_id is not None else 1
+        self.eos_token_id = self.tokenizer.eos_token_id if self.tokenizer.eos_token_id is not None else 2
+        self.assistant_output_token_id = 1565  # Token ID for "true", which is the last assistant's response in the example
path: tests/test_utils.py
text: @user1:
```suggestion
        # Token ID for "true", the last assistant's response in the example:
        self.assistant_output_token_id = 1565
```