Trainer does not properly read "label" column.
Author: notrichardren
Pull request: no

### System Info

When I have a HuggingFace dataset with 'input_ids', 'attention_mask', and 'labels', the current HuggingFace Trainer does not properly read 'labels' and separate the labels from the inputs in the forward pass.

I also tried a 'label' (as opposed to 'labels') column and it does not work with this either -- it fails to separate the label from the rest of the forward pass.

This contrasted with the demo for Trainer, where a dataset has the tokenized version of a given dataset as well as a 'label' feature column.

### Who can help?

@sgugger

### 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

The HuggingFace Accelerate config I'm using has fp16 enabled. Here is the code to replicate:

```
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)

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') # has columns 'input_ids', 'attention_mask', 'labels'
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

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

training_args = TrainingArguments("test-trainer", 
                                  num_train_epochs = 250,
                                  per_device_train_batch_size = 1,
                                  remove_unused_columns=False,
                                  fp16=True
)

trainer = Trainer( # probably using cross entropy loss
    reward_model,
    training_args,
    train_dataset=tokenized_dataset["train"],
    data_collator=data_collator,
    tokenizer=tokenizer,
)

trainer.train()

trainer.save_model("trained_model")
```

I receive a TypeError on the forward pass, where **forward() got an unexpected keyword argument 'labels'**
```
  File "minimum-viable-example.py", line 19, in forward
        return forward_call(*args, **kwargs)
  File "minimum-viable-example.py", line 19, in forward
tr_loss_step = self.training_step(model, inputs)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/transformers/trainer.py", line 2654, in training_step
    return inner_training_loop(
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/transformers/trainer.py", line 1809, in _inner_training_loop
    return forward_call(*args, **kwargs)
  File "minimum-viable-example.py", line 19, in forward
    outputs = model(**inputs)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    outputs = self.language_model(**args)    
ret_val = func(*args, **kwargs)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1735, in forward
    tr_loss_step = self.training_step(model, inputs)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/transformers/trainer.py", line 2654, in training_step
        outputs = self.language_model(**args)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
loss = self.compute_loss(model, inputs)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/transformers/trainer.py", line 2679, in compute_loss
    outputs = self.language_model(**args)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/torch/nn/modules/module.py", line 1501, in _call_impl
    loss = self.compute_loss(model, inputs)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/transformers/trainer.py", line 2679, in compute_loss
    return forward_call(*args, **kwargs)    
return forward_call(*args, **kwargs)
  File "/admin/home-notrichardren/.local/lib/python3.8/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
TypeError: forward() got an unexpected keyword argument 'labels'
    return forward_call(*args, **kwargs)
TypeError:     forward() got an unexpected keyword argument 'labels'loss = self.module(*inputs, **kwargs)
```

### Expected behavior

I would expect the 'labels' column (I tried both 'label' and 'labels') to not be put in the forward pass, but instead to be used in the loss function.

Comments:
You are using a model which does not accept `labels` in its forward pass: `AutoModel` gives you the base model, it is only suitable to extract the last hidden state generated by the model, not for training.

By default, am I supposed to pass a model to Trainer that is able to accept "labels" in its forward pass? 

I wasn't too sure how Trainer worked, but based on the documentation, my impression was that "labels" was supposed to be a separate column (that is then removed) before all of the args are passed into the forward pass. Afterward, the loss calculation would happen where the "labels" are compared to model predictions.

@sgugger 

@sgugger Thanks for the info. This error is spread over all of SO and other channels. But no one has pointed out that Automodel based derived models are not suitable for training. I faced this issue with T5Model and LlamaModel.
Can you point out the right way to select model for training?