Repository: huggingface/trl

Issue:
Exception handling in DataCollatorForCompletionOnlyLM
Hey,

While testing out the new `DataCollatorForCompletionOnlyLM` on custom data, I had some issues relating to the exception handling of batch instances in which we fail to find the `response_template` or `instruction_template` (https://github.com/huggingface/trl/blob/3b2c820db68aecca36c2ff5d152de59deec8fe5c/trl/trainer/utils.py#L142C3-L142C3). 

It turned out that I had some exceptionally long responses in my training data which didn't actually contain both a valid `response_template` and `instruction_template`.

**Possible fix:**

My current work around was to replace the `RuntimeError` with a more informative `warning` and simply set the labels for a problematic instance to `self.ignore_index`, so that it will be ignored for the loss calculation but training will still run successfully (provided this doesn't happen for every instance 😉 ), e.g.:

```python
if len(human_token_ids_idxs) == 0:
    warnings.warn(
        f'Could not find instruction key `{self.instruction_template}` in the following instance {self.tokenizer.decode(batch["input_ids"][i])} ' \
        f'This instance will be ignored in loss calculation.'
    )
    batch["labels"][i, :] = self.ignore_index
```

Would something like this be more suitable than the current exception handling?




Discussion:
Thanks for the investigation - that looks like a good fix to me! Would you like to open a PR? 

Resolving patch:
diff --git a/trl/trainer/utils.py b/trl/trainer/utils.py
--- a/trl/trainer/utils.py
+++ b/trl/trainer/utils.py
@@ -106,14 +106,18 @@ def torch_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> D
                         response_token_ids_start_idx = idx
 
                 if response_token_ids_start_idx is None:
-                    raise RuntimeError(
-                        f'Could not find response key {self.response_token_ids} in token IDs {batch["labels"][i]}'
+                    warnings.warn(
+                        f"Could not find response key `{self.response_template}` in the "
+                        f'following instance: {self.tokenizer.decode(batch["input_ids"][i])} '
+                        f"This instance will be ignored in loss calculation. "
+                        f"Note, if this happens often, consider increasing the `max_seq_length`."
                     )
+                    batch["labels"][i, :] = self.ignore_index
+                else:
+                    response_token_ids_end_idx = response_token_ids_start_idx + len(self.response_token_ids)
 
-                response_token_ids_end_idx = response_token_ids_start_idx + len(self.response_token_ids)
-
-                # Make pytorch loss function ignore all tokens up through the end of the response key
-                batch["labels"][i, :response_token_ids_end_idx] = self.ignore_index
+                    # Make pytorch loss function ignore all tokens up through the end of the response key
+                    batch["labels"][i, :response_token_ids_end_idx] = self.ignore_index
 
         else:
             for i in range(len(examples)):
@@ -128,10 +132,14 @@ def torch_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> D
                     ):
                         response_token_ids_idxs.append(assistant_idx + len(self.response_token_ids))
 
-                if len(self.response_token_ids) == 0:
-                    raise RuntimeError(
-                        f'Could not find response key {self.response_token_ids} in token IDs {batch["labels"][i]}'
+                if len(response_token_ids_idxs) == 0:
+                    warnings.warn(
+                        f"Could not find response key `{self.response_template}` in the "
+                        f'following instance: {self.tokenizer.decode(batch["input_ids"][i])} '
+                        f"This instance will be ignored in loss calculation. "
+                        f"Note, if this happens often, consider increasing the `max_seq_length`."
                     )
+                    batch["labels"][i, :] = self.ignore_index
 
                 human_token_ids = self.tokenizer.encode(self.instruction_template, add_special_tokens=False)
                 for human_idx in np.where(batch["labels"][i] == human_token_ids[0])[0]:
@@ -140,9 +148,13 @@ def torch_call(self, examples: List[Union[List[int], Any, Dict[str, Any]]]) -> D
                         human_token_ids_idxs.append(human_idx)
 
                 if len(human_token_ids_idxs) == 0:
-                    raise RuntimeError(
-                        f'Could not find response key {human_token_ids} in token IDs {batch["labels"][i]}'
+                    warnings.warn(
+                        f"Could not find instruction key `{self.instruction_template}` in the "
+                        f'following instance: {self.tokenizer.decode(batch["input_ids"][i])} '
+                        f"This instance will be ignored in loss calculation. "
+                        f"Note, if this happens often, consider increasing the `max_seq_length`."
                     )
+                    batch["labels"][i, :] = self.ignore_index
 
                 for idx, (start, end) in enumerate(zip(human_token_ids_idxs, response_token_ids_idxs)):
                     # Make pytorch loss function ignore all non response tokens


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
@@ -13,6 +13,7 @@
 # limitations under the License.
 import unittest
 
+import torch
 from transformers import AutoTokenizer
 
 from trl import DataCollatorForCompletionOnlyLM
@@ -20,16 +21,16 @@
 
 class DataCollatorForCompletionOnlyLMTester(unittest.TestCase):
     def test_data_collator_finds_response_template_llama2_tokenizer(self):
-        self.tokenizer = AutoTokenizer.from_pretrained(
-            "upstage/Llama-2-70b-instruct-v2"
-        )  # Not using official one to avoid logging in
+        self.tokenizer = AutoTokenizer.from_pretrained("trl-internal-testing/dummy-GPT2-correct-vocab")
         self.instruction = """### System: You are a helpful assistant.
 
 ### User: How much is 2+2?
 
 ### Assistant: 2+2 equals 4"""
         self.response_template = "\n### Assistant:"
-        # [29871, 13, 2277, 29937, 4007, 22137, 29901] -> [2277, 29937, 4007, 22137, 29901]
+
+        # GPT2Tokenizer: [198, 21017, 15286, 25] -> [15286, 25]
+        # Llama2Tokenizer: [29871, 13, 2277, 29937, 4007, 22137, 29901] -> [2277, 29937, 4007, 22137, 29901]
         self.tokenized_response_w_context = self.tokenizer.encode(self.response_template, add_special_tokens=False)[2:]
 
         # Plain check on string
@@ -40,3 +41,26 @@ def test_data_collator_finds_response_template_llama2_tokenizer(self):
         # Pass already tokenized (w context) and truncated response_template so token_ids are like in the instruction + response
         self.collator = DataCollatorForCompletionOnlyLM(self.tokenized_response_w_context, tokenizer=self.tokenizer)
         self.collator.torch_call([self.tokenized_instruction])
+
+    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.
+
+### User: How much is 2+2? I'm asking because I'm not sure. And I'm not sure because I'm not good at math.
+"""
+        self.response_template = "\n### Assistant:"
+        # check DataCollatorForCompletionOnlyLM using response template only
+        self.tokenized_instruction = self.tokenizer.encode(self.instruction, add_special_tokens=False)
+        self.collator = DataCollatorForCompletionOnlyLM(self.response_template, tokenizer=self.tokenizer)
+        encoded_instance = self.collator.torch_call([self.tokenized_instruction])
+        result = torch.all(encoded_instance["labels"] == -100)
+        self.assertTrue(result, "Not all values in the tensor are -100.")
+
+        # check DataCollatorForCompletionOnlyLM using response template and instruction template
+        self.instruction_template = "\n### User:"
+        self.collator = DataCollatorForCompletionOnlyLM(
+            self.response_template, self.instruction_template, tokenizer=self.tokenizer
+        )
+        encoded_instance = self.collator.torch_call([self.tokenized_instruction])
+        result = torch.all(encoded_instance["labels"] == -100)
+        self.assertTrue(result, "Not all values in the tensor are -100.")