Repository: huggingface/trl

Issue:
DataCollatorForCompletionOnlyLM instruction token masking fails if first occurence of instruction is marked differently
The current code in DataCollatorForCompletionOnlyLM assumes that the first deteced occurence of `instruction_template` comes before the first detected occurence of `response_template`. This is reasonable, since in current applications conversations are initiated by the user, not the assistant. However, this can fail if the first instruction is marked differently from all the other instructions, which can if a context-sensitive tokenizer such as Llama-2 tokenizes the instruction_template differently at the start of a string than in the middle.

In particular this happens in practice with TinyLlama: `
` gets tokenized as `529, 29989, 1792, 29989, 29958` at the start of a conversation, but as `29966, 29989, 1792, 29989, 29958` in later messages.


Reproduction snippet:
```
from transformers import AutoTokenizer
from trl import DataCollatorForCompletionOnlyLM

tokenizer = AutoTokenizer.from_pretrained("TinyLlama/TinyLlama-1.1B-Chat-v1.0")

chat = [
    {"role": "user", "content": "Which is bigger, the moon or the sun?"},
    {"role": "assistant", "content": "The sun."},
    {"role": "user", "content": "Really?"},
    {"role": "assistant", "content": "Yes."},
    {"role": "user", "content": "I don't believe you."},
    {"role": "assistant", "content": "That's okay."},
]

collator = DataCollatorForCompletionOnlyLM(response_template=[29966, 29989, 465, 22137, 29989, 29958, 13], instruction_template=[29871, 13, 29966, 29989, 1792, 29989, 29958], tokenizer=tokenizer)


collator([tokenizer.apply_chat_template(chat)])
# This doesn't mask user messages 2 & 3.
```

PR #1185 fixes this, and makes the above snippet mask out all the user messages correctly.


Discussion:

Resolving patch:
diff --git a/trl/trainer/utils.py b/trl/trainer/utils.py
--- a/trl/trainer/utils.py
+++ b/trl/trainer/utils.py
@@ -176,6 +176,13 @@ def torch_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> D
                     )
                     batch["labels"][i, :] = self.ignore_index
 
+                if (
+                    len(human_token_ids_idxs) > 0
+                    and len(response_token_ids_idxs) > 0
+                    and human_token_ids_idxs[0] > response_token_ids_idxs[0]
+                ):
+                    human_token_ids_idxs = [0] + human_token_ids_idxs
+
                 for idx, (start, end) in enumerate(zip(human_token_ids_idxs, response_token_ids_idxs)):
                     # Make pytorch loss function ignore all non response tokens
                     if idx != 0:


Test patch:
diff --git a/tests/test_data_collator_completion_only.py b/tests/test_data_collator_completion_only.py
--- a/tests/test_data_collator_completion_only.py
+++ b/tests/test_data_collator_completion_only.py
@@ -31,11 +31,14 @@ def test_data_collator_finds_response_template_llama2_tokenizer(self):
         self.instruction_template = "\n### User:"
         self.response_template = "\n### Assistant:"
 
-        # GPT2Tokenizer: [198, 21017, 11787, 25] -> [11787, 25]
+        # GPT2Tokenizer: [198, 21017, 11787, 25] -> [21017, 11787, 25]
         # Llama2Tokenizer: [29871, 13, 2277, 29937, 4911, 29901] -> [2277, 29937, 4911, 29901]
+        # Note: If this test is ever switched to Llama2Tokenizer, this should be double checked,
+        # and possibly switched back to [2:] instead of [1:].
+        # With GPT2Tokenizer, [1:] is correct - we want the 21017 token included, which is ###.
         self.tokenized_instruction_w_context = self.tokenizer.encode(
             self.instruction_template, add_special_tokens=False
-        )[2:]
+        )[1:]
 
         # GPT2Tokenizer: [198, 21017, 15286, 25] -> [15286, 25]
         # Llama2Tokenizer: [29871, 13, 2277, 29937, 4007, 22137, 29901] -> [2277, 29937, 4007, 22137, 29901]
@@ -57,6 +60,28 @@ def test_data_collator_finds_response_template_llama2_tokenizer(self):
         )
         self.collator.torch_call([self.tokenized_instruction])
 
+        # Test for PR #1185
+        # We pass in a string where the first user template is different than the rest.
+        # Usually this would happen due to context-sensitive tokenization, but here we
+        # explicitly change the template to test the fix.
+        self.instruction = """## User: First instruction
+
+### Assistant: First response
+
+### User: Second instruction
+
+### Assistant: Second response"""
+        self.tokenized_instruction = self.tokenizer.encode(self.instruction, add_special_tokens=False)
+        self.collator = DataCollatorForCompletionOnlyLM(
+            self.tokenized_response_w_context, self.tokenized_instruction_w_context, tokenizer=self.tokenizer
+        )
+        collator_output = self.collator.torch_call([self.tokenized_instruction])
+        collator_text = self.tokenizer.decode(
+            collator_output["labels"][torch.where(collator_output["labels"] != -100)]
+        )
+        expected_text = " First response\n\n Second response" ""
+        self.assertEqual(collator_text, expected_text)
+
     def test_data_collator_handling_of_long_sequences(self):
         self.tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/dummy-GPT2-correct-vocab")
         self.instruction = """### System: You are a helpful assistant.