removed_unused_columns = True deletes dataset for custom loss func
Author: notrichardren
Pull request: no

Feel free to see the[ second comment](https://github.com/huggingface/transformers/issues/25054#issuecomment-1648263249) directly. The original issue was on confusion for the training dataset seemingly not working -- I narrowed it down to the ```removed_unused_columns = True``` dataset deleting an entire dataset, which contradicts the description in the documentation.

### System Info

- `transformers` version: 4.31.0
- Platform: Linux-5.15.0-1037-aws-x86_64-with-glibc2.29
- Python version: 3.8.10
- Huggingface_hub version: 0.15.1
- Safetensors version: 0.3.1
- Accelerate version: 0.21.0
- Accelerate config:    - compute_environment: LOCAL_MACHINE
        - distributed_type: DEEPSPEED
        - mixed_precision: fp16
        - use_cpu: False
        - num_processes: 7
        - machine_rank: 0
        - num_machines: 1
        - rdzv_backend: static
        - same_network: True
        - main_training_function: main
        - deepspeed_config: {'gradient_accumulation_steps': 1, 'offload_optimizer_device': 'none', 'offload_param_device': 'none', 'zero3_init_flag': False, 'zero_stage': 2}
        - downcast_bf16: no
        - tpu_use_cluster: False
        - tpu_use_sudo: False
        - tpu_env: []
- PyTorch version (GPU?): 2.0.1+cu118 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: GPU
- Using distributed or parallel set-up in script?: parallel

### Who can help?

@sgugger

I sincerely believe this is a bug in Transformers, though my apologies if it turns out to be my own code.

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

```
import torch
from torch import nn
from transformers import AutoTokenizer, AutoModel, LlamaTokenizer, DataCollatorWithPadding
import pandas as pd
from datasets import load_dataset, Dataset, DatasetDict
from functools import partial
from transformers import TrainingArguments, Trainer
import numpy as np
import evaluate

class RewardModel(nn.Module):

    def __init__(self, model):
        super().__init__()
        self.language_model = model
        self.fc = nn.Linear(self.language_model.config.hidden_size, 1)

    def forward(self, **args):
        outputs = self.language_model(**args)
        last_hidden_state = outputs.last_hidden_state
        reward = self.fc(last_hidden_state) # (batch_size, seq_len, 1)
        reward = reward.squeeze(-1) # (batch_size, seq_len)
        reward = reward[:,-1] # takes reward at last seq pos (batch_size)
        return reward

pretrained_model_name = "decapoda-research/llama-7b-hf"
model = AutoModel.from_pretrained(pretrained_model_name)
reward_model = RewardModel(model)

for param in reward_model.parameters(): # all the requires grads are false
    param.requires_grad = False
for param in reward_model.fc.parameters(): # except the last layer
    param.requires_grad = True

tokenizer = LlamaTokenizer.from_pretrained(pretrained_model_name)
if tokenizer.pad_token is None:
    tokenizer.pad_token='[PAD]'

tokenized_dataset = load_dataset('notrichardren/hh-rlhf-tf') # datasetdict with train/test split that has columns 'input_ids', 'attention_mask', 'labels'
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

class LabelPopTrainer(Trainer):
    def compute_loss(self,model,inputs, return_outputs=False):
        labels = inputs.pop("labels")
        outputs = model(**inputs).flatten()
        loss = torch.nn.functional.cross_entropy(outputs, labels.half())
        return (loss, outputs) if return_outputs else loss

args = TrainingArguments("test-trainer", 
                        num_train_epochs = 3,
                        per_device_train_batch_size = 4,
                        logging_strategy = "steps",
                        logging_steps = 3,
                        fp16=True
)

trainer = LabelPopTrainer(
    reward_model,
    args,
    train_dataset=tokenized_dataset["train"],
    data_collator=data_collator,
    tokenizer=tokenizer,
)

trainer.train()

trainer.save_model("trained_model")
```

### Expected behavior

I get the error that “Invalid key: 222838 is out of bounds for size 0” which is due to the train_dataloader __getitem__ function not working. However, tokenized_dataset[“train”] (which is the one passed to the trainer) is:

```
Dataset({
    features: ['labels', 'input_ids', 'attention_mask'],
    num_rows: 224104
})
```

I would expect the dataset, therefore, to work when run with ```accelerate launch ___.py``` Below is the more complete error:

```
Traceback (most recent call last):
  File "reward-model-v3-mve-2.py", line 67, in <module>
    trainer.train()
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/transformers/trainer.py", line 1539, in train
    return inner_training_loop(
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/transformers/trainer.py", line 1787, in _inner_training_loop
    for step, inputs in enumerate(epoch_iterator):
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/accelerate/data_loader.py", line 384, in __iter__
    current_batch = next(dataloader_iter)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 633, in __next__
    data = self._next_data()
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 677, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 49, in fetch
    data = self.dataset.__getitems__(possibly_batched_index)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 2796, in __getitems__
    batch = self.__getitem__(keys)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 2792, in __getitem__
    return self._getitem(key)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/datasets/arrow_dataset.py", line 2776, in _getitem
    pa_subtable = query_table(self._data, key, indices=self._indices if self._indices is not None else None)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/datasets/formatting/formatting.py", line 583, in query_table
    _check_valid_index_key(key, size)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/datasets/formatting/formatting.py", line 536, in _check_valid_index_key
    _check_valid_index_key(int(max(key)), size=size)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-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: 202710 is out of bounds for size 0
```

Comments:
I did some further investigation and found that this dataset seems to work _if_ and only if you set remove_unused_columns = False in the training arguments.

**It seems very odd that "remove_unused_columns = True" would result in the entire dataset being wiped, given that the features in the dataset are ['labels', 'input_ids', 'attention_mask'].** This on its own seems like a bug rather than a feature and seems like it could be a point of confusion (especially without an appropriate error message for the user). It seems to also directly contradict the [documentation description](https://huggingface.co/docs/transformers/main_classes/trainer#transformers.TrainingArguments.remove_unused_columns) of the function.

If you'd like to modify remove_unused_columns to reflect this, it may be good to keep the issue open to address this -- otherwise, feel free to close it.

`remove_unused_columns=True` will remove the keys in the dataset that are not accepted by the models, which is decided by looking at the model forward signature. Your `RewardModel` does not have any argument names (it takes `*args`), that's why there is this bug. If you unpack that `*args` to use names like `labels`, `input_ids`, `attention_mask`, the issue will disappear.

Ah, I see. Thank you very much for clarifying this, this makes a lot of sense. I appreciate it!

No problem! We should probably make a note of it in the documentation.