Repository: huggingface/trl
Pull request: [`core` / SFTTrainer] Fix breaking change
# What does this PR do?

fixes a breaking change on a recent PR: https://github.com/huggingface/trl/pull/1188 
Fixes: https://github.com/huggingface/trl/issues/1216 

In fact, if one uses the default data collator + `remove_unused_columns=False`, the training script fails for some datasets such as imdb as the collator will try to encode all the columns, including `"text"`. The fix should be to force-set remove_unused_columns=True, only in the case users do not pass a data_collator for non-packing datasets and explain to the users who they should proceed in case they decided to pass `remove_unused_columns=False`

cc @lvwerra 

Problem statement:
Failed to load data in trl 0.7.8/0.7.9.
This is a new regression introduced in trl 0.7.8 (and 0.7.9), 0.7.7 is fine.

We run into issues of `ValueError: too many dimensions 'str'` when loading data to the trainer.
Here's a simple LLAMA2+LoRA fine-tuning on IMDB dataset as minimal repro:

```py3
#!/usr/bin/env python3

import datasets
import peft
import transformers
import trl


model_dir = "models/Llama-2-7b-hf"

tokenizer = transformers.AutoTokenizer.from_pretrained(model_dir)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

model = transformers.AutoModelForCausalLM.from_pretrained(model_dir)

ds_train = datasets.load_dataset("imdb", split="train[:10]")

trainer = trl.SFTTrainer(
    model=model,
    args=transformers.TrainingArguments(
        output_dir="output",
        max_steps=1,
        remove_unused_columns=False,
    ),
    peft_config=peft.LoraConfig(
        lora_alpha=16,
        lora_dropout=0.1,
        r=8,
        bias="none",
        task_type="Causal_LM",
    ),
    train_dataset=ds_train,
    tokenizer=tokenizer,
    dataset_text_field="text",
    max_seq_length=8,
)
trainer.train()
```

0.7.7 works:
```py3
# CUDA_VISIBLE_DEVICES=0 ./test.py 
/usr/lib/python3/dist-packages/requests/__init__.py:87: RequestsDependencyWarning: urllib3 (2.1.0) or chardet (4.0.0) doesn't match a supported version!
  warnings.warn("urllib3 ({}) or chardet ({}) doesn't match a supported "
Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:03<00:00,  1.52s/it]
Detected kernel version 3.10.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.
  0%|                                                                                                                                                                                                                                 | 0/1 [00:00<?, ?it/s]You're using a LlamaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  1.24it/s]Attempted to log scalar metric train_runtime:
0.8097
Attempted to log scalar metric train_samples_per_second:
9.88
Attempted to log scalar metric train_steps_per_second:
1.235
Attempted to log scalar metric total_flos:
2538830561280.0
Attempted to log scalar metric train_loss:
4.124451637268066
Attempted to log scalar metric epoch:
0.5
{'train_runtime': 0.8097, 'train_samples_per_second': 9.88, 'train_steps_per_second': 1.235, 'train_loss': 4.124451637268066, 'epoch': 0.5}                                                                                                                 
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1/1 [00:00<00:00,  1.24it/s]
```

0.7.8 failed:
```py3
# CUDA_VISIBLE_DEVICES=0 ./test.py 
/usr/lib/python3/dist-packages/requests/__init__.py:87: RequestsDependencyWarning: urllib3 (2.1.0) or chardet (4.0.0) doesn't match a supported version!
  warnings.warn("urllib3 ({}) or chardet ({}) doesn't match a supported "
Loading checkpoint shards: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [00:03<00:00,  1.54s/it]
Map: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 940.17 examples/s]
Detected kernel version 3.10.0, which is below the recommended minimum of 5.5.0; this can cause the process to hang. It is recommended to upgrade the kernel to the minimum version or higher.
  0%|                                                                                                                                                                                                                                 | 0/1 [00:00<?, ?it/s]You're using a LlamaTokenizerFast tokenizer. Please note that with a fast tokenizer, using the `__call__` method is faster than using a method to encode the text followed by a call to the `pad` method to get a padded encoding.
Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/transformers/tokenization_utils_base.py", line 748, in convert_to_tensors
    tensor = as_tensor(value)
  File "/usr/local/lib/python3.10/dist-packages/transformers/tokenization_utils_base.py", line 720, in as_tensor
    return torch.tensor(value)
ValueError: too many dimensions 'str'

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "./test.py", line 38, in <module>
    trainer.train()
  File "/usr/local/lib/python3.10/dist-packages/trl/trainer/sft_trainer.py", line 317, in train
    output = super().train(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 1537, in train
    return inner_training_loop(
  File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 1821, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "/usr/local/lib/python3.10/dist-packages/accelerate/data_loader.py", line 448, in __iter__
    current_batch = next(dataloader_iter)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py", line 630, in __next__
    data = self._next_data()
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py", line 674, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/fetch.py", line 54, in fetch
    return self.collate_fn(data)
  File "/usr/local/lib/python3.10/dist-packages/transformers/data/data_collator.py", line 45, in __call__
    return self.torch_call(features)
  File "/usr/local/lib/python3.10/dist-packages/transformers/data/data_collator.py", line 732, in torch_call
    batch = self.tokenizer.pad(examples, return_tensors="pt", pad_to_multiple_of=self.pad_to_multiple_of)
  File "/usr/local/lib/python3.10/dist-packages/transformers/tokenization_utils_base.py", line 3299, in pad
    return BatchEncoding(batch_outputs, tensor_type=return_tensors)
  File "/usr/local/lib/python3.10/dist-packages/transformers/tokenization_utils_base.py", line 223, in __init__
    self.convert_to_tensors(tensor_type=tensor_type, prepend_batch_axis=prepend_batch_axis)
  File "/usr/local/lib/python3.10/dist-packages/transformers/tokenization_utils_base.py", line 764, in convert_to_tensors
    raise ValueError(
ValueError: Unable to create tensor, you should probably activate truncation and/or padding with 'padding=True' 'truncation=True' to have batched tensors with the same length. Perhaps your features (`text` in this case) have excessive nesting (inputs type `list` where type `int` is expected).
  0%|          | 0/1 [00:00<?, ?it/s]
```

Hints from issue discussion:
Thanks for the repro @xkszltl , will try to repro and fix the issue and make another patch release
Hi @xkszltl 
it seems to be caused by `remove_unused_columns=False,`, can you meanwhile either revert to 0.7.7 or set `remove_unused_columns=True` ? I'll try to provide the right fix meanwhile
True won't work at all (regardless of the version), that's why it's False initially.
My understanding is it drops the "text" column because model only has "input_ids" on its interface?

```py3
Traceback (most recent call last):
  File "./test.py", line 58, in main
    trainer.train()
  File "/usr/local/lib/python3.10/dist-packages/trl/trainer/sft_trainer.py", line 315, in train
    output = super().train(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 1537, in train
    return inner_training_loop(
  File "/usr/local/lib/python3.10/dist-packages/transformers/trainer.py", line 1821, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "/usr/local/lib/python3.10/dist-packages/accelerate/data_loader.py", line 451, in __iter__
    current_batch = next(dataloader_iter)
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py", line 630, in __next__
    data = self._next_data()
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py", line 674, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = self.dataset.__getitems__(possibly_batched_index)
  File "/usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py", line 2804, in __getitems__
    batch = self.__getitem__(keys)
  File "/usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py", line 2800, in __getitem__
    return self._getitem(key)
  File "/usr/local/lib/python3.10/dist-packages/datasets/arrow_dataset.py", line 2784, in _getitem
    pa_subtable = query_table(self._data, key, indices=self._indices if self._indices is not None else None)
  File "/usr/local/lib/python3.10/dist-packages/datasets/formatting/formatting.py", line 583, in query_table
    _check_valid_index_key(key, size)
  File "/usr/local/lib/python3.10/dist-packages/datasets/formatting/formatting.py", line 536, in _check_valid_index_key
    _check_valid_index_key(int(max(key)), size=size)
  File "/usr/local/lib/python3.10/dist-packages/datasets/formatting/formatting.py", line 526, in _check_valid_index_key
    raise IndexError(f"Invalid key: {key} is out of bounds for size {size}")
IndexError: Invalid key: 23887 is out of bounds for size 0
```
I'm currently pinning to `trl<0.7.8`.
hi @xkszltl 
Thnks for your patience! I had a deeper look at the issue and I made https://github.com/huggingface/trl/pull/1229 that should resolve it.
Regarding https://github.com/huggingface/trl/issues/1216#issuecomment-1886387405 - can you try to update datasets? I cannot repro with:

```python
import datasets
import peft
import transformers
import trl


model_dir = "HuggingFaceM4/tiny-random-LlamaForCausalLM"

tokenizer = transformers.AutoTokenizer.from_pretrained(model_dir)
tokenizer.pad_token = tokenizer.eos_token
tokenizer.padding_side = "right"

model = transformers.AutoModelForCausalLM.from_pretrained(model_dir)

ds_train = datasets.load_dataset("imdb", split="train[:10]")

trainer = trl.SFTTrainer(
    model=model,
    args=transformers.TrainingArguments(
        output_dir="output",
        max_steps=1,
        remove_unused_columns=True,
    ),
    peft_config=peft.LoraConfig(
        lora_alpha=16,
        lora_dropout=0.1,
        r=8,
        bias="none",
        task_type="Causal_LM",
    ),
    train_dataset=ds_train,
    tokenizer=tokenizer,
    dataset_text_field="text",
    max_seq_length=8,
)
trainer.train()
```
Are you trying with master or a release? 
@xkszltl on master currently
@xkszltl can you try and let me know how it goes?
Only tried on released wheel so that may be the reason.
I can give master a try after that PR is merged. 
I see ok ! if you want you can build from that branch:
```bash
pip install -U git+https://github.com/huggingface/trl.git@fix-breaking-change
```

Commit message: fix

Patch to review:
diff --git a/trl/trainer/sft_trainer.py b/trl/trainer/sft_trainer.py
index 8d3320969e..3743f299cd 100644
--- a/trl/trainer/sft_trainer.py
+++ b/trl/trainer/sft_trainer.py
@@ -243,6 +243,10 @@ def make_inputs_require_grad(module, input, output):
             # if not stays #None
             formatting_func = get_formatting_func_from_dataset(train_dataset, tokenizer)
 
+        # This variable is used to check later on if one has loaded `SFTtrainer` with the
+        # default collator or not.
+        use_default_collator = True
+
         if not packing:
             if dataset_text_field is None and formatting_func is None:
                 raise ValueError(
@@ -251,6 +255,8 @@ def make_inputs_require_grad(module, input, output):
 
             if data_collator is None:
                 data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
+            else:
+                use_default_collator = False
 
         if dataset_kwargs is None:
             dataset_kwargs = {}
@@ -265,6 +271,7 @@ def make_inputs_require_grad(module, input, output):
                 num_of_sequences,
                 chars_per_token,
                 remove_unused_columns=args.remove_unused_columns if args is not None else True,
+                use_default_collator=use_default_collator,
                 **dataset_kwargs,
             )
         if eval_dataset is not None:
@@ -359,6 +366,7 @@ def _prepare_dataset(
         remove_unused_columns=True,
         append_concat_token=True,
         add_special_tokens=True,
+        use_default_collator=True,
     ):
         if dataset is None:
             raise ValueError("The dataset should not be None")
@@ -376,6 +384,7 @@ def _prepare_dataset(
                 formatting_func,
                 add_special_tokens,
                 remove_unused_columns,
+                use_default_collator,
             )
 
         else:
@@ -400,6 +409,7 @@ def _prepare_non_packed_dataloader(
         formatting_func=None,
         add_special_tokens=True,
         remove_unused_columns=True,
+        use_default_collator=True,
     ):
         use_formatting_func = formatting_func is not None and dataset_text_field is None
         self._dataset_sanity_checked = False
@@ -424,7 +434,18 @@ def tokenize(element):
                 else:
                     self._dataset_sanity_checked = True
 
-            return {"input_ids": outputs["input_ids"], "attention_mask": outputs["attention_mask"]}
+            return {
+                "input_ids": outputs["input_ids"],
+                "labels": outputs["input_ids"],
+                "attention_mask": outputs["attention_mask"],
+            }
+
+        if not remove_unused_columns and use_default_collator:
+            warnings.warn(
+                "You passed `remove_unused_columns=False` on a non-packed dataset while using `DataCollatorForLanguageModeling`. This might create some issues with the default collator. If you want to "
+                "inspect dataset other columns, you can subclass `DataCollatorForLanguageModeling` and create your own data collator in order to inspect the unused dataset columns."
+            )
+            remove_unused_columns = True
 
         tokenized_dataset = dataset.map(
             tokenize,


Code review comments:
diff_hunk: @@ -424,7 +434,18 @@ def tokenize(element):
                 else:
                     self._dataset_sanity_checked = True
 
-            return {"input_ids": outputs["input_ids"], "attention_mask": outputs["attention_mask"]}
+            return {
+                "input_ids": outputs["input_ids"],
+                "labels": outputs["input_ids"],
+                "attention_mask": outputs["attention_mask"],
+            }
+
+        if not remove_unused_columns and use_default_collator:
+            warnings.warn(
+                "You passed `remove_unused_columns=False` on a non-packed dataset while using `DataCollatorForLanguageModeling`. This might create some issues with the default collator. If you want to "
+                "inspect dataset other columns, you can subclass `DataCollatorForLanguageModeling` and create your own data collator in order to inspect the unused dataset columns."
+            )
+            remove_unused_columns = True
path: trl/trainer/sft_trainer.py
text: @user1:
I still think we shouldn't overwrite if the user provided `remove_unused_colums=False` and not always yield a warning because there might not be an issue at all.