Trainer misshaping input_ids
Author: lhallee
Pull request: no

### System Info

2023-11-10 17:12:03.551843: E tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:9342] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
2023-11-10 17:12:03.551902: E tensorflow/compiler/xla/stream_executor/cuda/cuda_fft.cc:609] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
2023-11-10 17:12:03.551946: E tensorflow/compiler/xla/stream_executor/cuda/cuda_blas.cc:1518] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
2023-11-10 17:12:04.651349: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT
WARNING:tensorflow:From /usr/local/lib/python3.10/dist-packages/transformers/commands/env.py:100: is_gpu_available (from tensorflow.python.framework.test_util) is deprecated and will be removed in a future version.
Instructions for updating:
Use `tf.config.list_physical_devices('GPU')` instead.
2023-11-10 17:12:07.781203: W tensorflow/core/common_runtime/gpu/gpu_bfc_allocator.cc:47] Overriding orig_value setting because the TF_FORCE_GPU_ALLOW_GROWTH environment variable is set. Original config value was 0.

Copy-and-paste the text below in your GitHub issue and FILL OUT the two last points.

- `transformers` version: 4.35.0
- Platform: Linux-5.15.120+-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.24.1
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.1.0+cu118 (True)
- Tensorflow version (GPU?): 2.14.0 (True)
- Flax version (CPU?/GPU?/TPU?): 0.7.5 (gpu)
- Jax version: 0.4.20
- JaxLib version: 0.4.20
- Using GPU in script?: <fill in>
- Using distributed or parallel set-up in script?: <fill in>

### Who can help?

@muellerzr @pacman100 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

https://huggingface.co/transformers/v3.1.0/custom_datasets.html
The dataset as portrayed leads to flattened inputs that are 1d batch_size * seq_len
Returning tensors pt also leads to batch_size, 1, seq_len which is also wrong.

```
class FineTuneDatasetIDS(Dataset):
    def __init__(self, seqs, labels, tokenizer, max_length=512):
        self.seqs = seqs
        self.labels = labels
        self.tokenizer = tokenizer
        self.max_length = max_length

    def __getitem__(self, idx):
        seq = self.seqs[idx]
        labels = self.labels[idx]
        encodings = self.tokenizer(seq,
                                   add_special_tokens=True,
                                   padding='max_length',
                                   max_length=self.max_length,
                                   truncation=True)
        item = {key: torch.tensor(val) for key, val in encodings.items()}
        item['labels'] = torch.tensor(labels, dtype=torch.float)
        return item

    def __len__(self):
        return len(self.labels)

training_args = TrainingArguments(
    output_dir=f'./results_{experiment}',
    num_train_epochs=100,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    warmup_steps=1000,
    weight_decay=0.01,
    do_train=True,
    do_eval=True,
    logging_dir=f'./logs_{experiment}',
    logging_steps=100,
    learning_rate=2e-5,
    adam_beta1 = 0.9,
    adam_beta2 = 0.98,
    evaluation_strategy='epoch',
    gradient_accumulation_steps=16,
    fp16=False,
    fp16_opt_level='02',
    run_name=experiment,
    seed=42,
    load_best_model_at_end=True,
    metric_for_best_model='fmax',
    greater_is_better=True,
    save_strategy='epoch'
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=valid_dataset,
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=5)]
)

trainer.train()
predictions, labels, metrics_output = trainer.predict(test_dataset)
metrics_output
```


```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
[<ipython-input-19-a560f54159b5>](https://localhost:8080/#) in <cell line: 1>()
----> 1 trainer.train()
      2 predictions, labels, metrics_output = trainer.predict(test_dataset)
      3 metrics_output

13 frames
[/usr/local/lib/python3.10/dist-packages/transformers/trainer.py](https://localhost:8080/#) in train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1553                 hf_hub_utils.enable_progress_bars()
   1554         else:
-> 1555             return inner_training_loop(
   1556                 args=args,
   1557                 resume_from_checkpoint=resume_from_checkpoint,

[/usr/local/lib/python3.10/dist-packages/transformers/trainer.py](https://localhost:8080/#) in _inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   1858 
   1859                 with self.accelerator.accumulate(model):
-> 1860                     tr_loss_step = self.training_step(model, inputs)
   1861 
   1862                 if (

[/usr/local/lib/python3.10/dist-packages/transformers/trainer.py](https://localhost:8080/#) in training_step(self, model, inputs)
   2723 
   2724         with self.compute_loss_context_manager():
-> 2725             loss = self.compute_loss(model, inputs)
   2726 
   2727         if self.args.n_gpu > 1:

[/usr/local/lib/python3.10/dist-packages/transformers/trainer.py](https://localhost:8080/#) in compute_loss(self, model, inputs, return_outputs)
   2746         else:
   2747             labels = None
-> 2748         outputs = model(**inputs)
   2749         # Save past state if it exists
   2750         # TODO: this needs to be fixed and made cleaner later.

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1516             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517         else:
-> 1518             return self._call_impl(*args, **kwargs)
   1519 
   1520     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1525                 or _global_backward_pre_hooks or _global_backward_hooks
   1526                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527             return forward_call(*args, **kwargs)
   1528 
   1529         try:

[<ipython-input-7-6912ec56d6d8>](https://localhost:8080/#) in forward(self, input_ids, attention_mask, labels)
     57 
     58     def forward(self, input_ids, attention_mask, labels=None):
---> 59         hidden = self.plm(input_ids=input_ids, attention_mask=attention_mask, labels=labels).hidden_states[-1]
     60         logits = self.classifier(hidden[:, 0, :])
     61         if labels is not None:

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1516             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517         else:
-> 1518             return self._call_impl(*args, **kwargs)
   1519 
   1520     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1525                 or _global_backward_pre_hooks or _global_backward_hooks
   1526                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527             return forward_call(*args, **kwargs)
   1528 
   1529         try:

[/usr/local/lib/python3.10/dist-packages/transformers/models/esm/modeling_esm.py](https://localhost:8080/#) in forward(self, input_ids, attention_mask, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, labels, output_attentions, output_hidden_states, return_dict)
   1026 
   1027             labels = labels.to(prediction_scores.device)
-> 1028             masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
   1029 
   1030         if not return_dict:

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1516             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517         else:
-> 1518             return self._call_impl(*args, **kwargs)
   1519 
   1520     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1525                 or _global_backward_pre_hooks or _global_backward_hooks
   1526                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527             return forward_call(*args, **kwargs)
   1528 
   1529         try:

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/loss.py](https://localhost:8080/#) in forward(self, input, target)
   1177 
   1178     def forward(self, input: Tensor, target: Tensor) -> Tensor:
-> 1179         return F.cross_entropy(input, target, weight=self.weight,
   1180                                ignore_index=self.ignore_index, reduction=self.reduction,
   1181                                label_smoothing=self.label_smoothing)

[/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py](https://localhost:8080/#) in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
   3051     if size_average is not None or reduce is not None:
   3052         reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 3053     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
   3054 
   3055 

ValueError: Expected input batch_size (1024) to match target batch_size (1170).
```



### Expected behavior

The trainer compiles a dataloader with the dataset that returns batch_size, seq_len entires.

Comments:
The updated example gives the same issue
https://huggingface.co/docs/transformers/tasks/sequence_classification

```
class ClassificationHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.out_proj = nn.Linear(config.hidden_size, config.num_labels)

    def forward(self, features, **kwargs):
        x = features[:, 0, :]
        x = self.dropout(x)
        x = self.dense(x)
        x = torch.tanh(x)
        x = self.dropout(x)
        x = self.out_proj(x)
        return x

class SequenceClassifier(nn.Module):
    def __init__(self, model_path, num_labels, weight_path=None, class_type='multilabel'):
        super(SequenceClassifier, self).__init__()
        self.config = AutoConfig.from_pretrained(model_path)
        self.config.num_labels = num_labels
        self.config.output_hidden_states = True
        self.plm = AutoModelForMaskedLM.from_pretrained(model_path, config=self.config)
        self.classifier = ClassificationHead(self.config)
        if weight_path is not None:
            self.plm.load_state_dict(torch.load(weight_path))
        if class_type == 'singlelabel':
            self.loss_fct = nn.CrossEntropyLoss()
        elif class_type == 'multilabel':
            self.loss_fct = nn.BCEWithLogitsLoss()
        self.num_labels = num_labels

    def forward(self, input_ids, attention_mask, labels=None):
        hidden = self.plm(input_ids=input_ids, attention_mask=attention_mask, labels=labels).hidden_states[-1]
        logits = self.classifier(hidden[:, 0, :])
        if labels is not None:
            loss = self.loss_fct(logits.view(-1, self.num_labels), labels.view(-1, self.num_labels))
        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=None,
            attentions=None
        )

def preprocess_function(examples):
    return tokenizer(examples['sequences'], truncation=True, max_length=512)

train_dataset = Dataset.from_dict({'sequences': train_seqs, 'labels': train_labels})
valid_dataset = Dataset.from_dict({'sequences': valid_seqs, 'labels': valid_labels})
test_dataset = Dataset.from_dict({'sequences': test_seqs, 'labels': test_labels})

dataset_dict = DatasetDict({
    'train': train_dataset,
    'validation': valid_dataset,
    'test': test_dataset
})

tokenized_data = dataset_dict.map(preprocess_function, batched=True)
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)
model = SequenceClassifier(model_path, len(train_labels[0]), weight_path, 'multilabel')

training_args = TrainingArguments(
    output_dir=f'./results_{experiment}',
    num_train_epochs=100,
    per_device_train_batch_size=2,
    per_device_eval_batch_size=2,
    warmup_steps=1000,
    weight_decay=0.01,
    do_train=True,
    do_eval=True,
    logging_dir=f'./logs_{experiment}',
    logging_steps=100,
    learning_rate=2e-5,
    adam_beta1 = 0.9,
    adam_beta2 = 0.98,
    evaluation_strategy='epoch',
    gradient_accumulation_steps=16,
    fp16=False,
    fp16_opt_level='02',
    run_name=experiment,
    seed=42,
    load_best_model_at_end=True,
    metric_for_best_model='fmax',
    greater_is_better=True,
    save_strategy='epoch'
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_data['train'],
    eval_dataset=tokenized_data['validation'],
    data_collator=data_collator,
    compute_metrics=compute_metrics,
    callbacks=[EarlyStoppingCallback(early_stopping_patience=5)]
)

trainer.train()
predictions, labels, metrics_output = trainer.predict(test_dataset)
metrics_output
```

```
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
[<ipython-input-18-a560f54159b5>](https://localhost:8080/#) in <cell line: 1>()
----> 1 trainer.train()
      2 predictions, labels, metrics_output = trainer.predict(test_dataset)
      3 metrics_output

13 frames
[/usr/local/lib/python3.10/dist-packages/transformers/trainer.py](https://localhost:8080/#) in train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1553                 hf_hub_utils.enable_progress_bars()
   1554         else:
-> 1555             return inner_training_loop(
   1556                 args=args,
   1557                 resume_from_checkpoint=resume_from_checkpoint,

[/usr/local/lib/python3.10/dist-packages/transformers/trainer.py](https://localhost:8080/#) in _inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   1858 
   1859                 with self.accelerator.accumulate(model):
-> 1860                     tr_loss_step = self.training_step(model, inputs)
   1861 
   1862                 if (

[/usr/local/lib/python3.10/dist-packages/transformers/trainer.py](https://localhost:8080/#) in training_step(self, model, inputs)
   2723 
   2724         with self.compute_loss_context_manager():
-> 2725             loss = self.compute_loss(model, inputs)
   2726 
   2727         if self.args.n_gpu > 1:

[/usr/local/lib/python3.10/dist-packages/transformers/trainer.py](https://localhost:8080/#) in compute_loss(self, model, inputs, return_outputs)
   2746         else:
   2747             labels = None
-> 2748         outputs = model(**inputs)
   2749         # Save past state if it exists
   2750         # TODO: this needs to be fixed and made cleaner later.

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1516             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517         else:
-> 1518             return self._call_impl(*args, **kwargs)
   1519 
   1520     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1525                 or _global_backward_pre_hooks or _global_backward_hooks
   1526                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527             return forward_call(*args, **kwargs)
   1528 
   1529         try:

[<ipython-input-13-cf6ab19ee4f9>](https://localhost:8080/#) in forward(self, input_ids, attention_mask, labels)
     59 
     60     def forward(self, input_ids, attention_mask, labels=None):
---> 61         hidden = self.plm(input_ids=input_ids, attention_mask=attention_mask, labels=labels).hidden_states[-1]
     62         logits = self.classifier(hidden[:, 0, :])
     63         if labels is not None:

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1516             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517         else:
-> 1518             return self._call_impl(*args, **kwargs)
   1519 
   1520     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1525                 or _global_backward_pre_hooks or _global_backward_hooks
   1526                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527             return forward_call(*args, **kwargs)
   1528 
   1529         try:

[/usr/local/lib/python3.10/dist-packages/transformers/models/esm/modeling_esm.py](https://localhost:8080/#) in forward(self, input_ids, attention_mask, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, labels, output_attentions, output_hidden_states, return_dict)
   1026 
   1027             labels = labels.to(prediction_scores.device)
-> 1028             masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
   1029 
   1030         if not return_dict:

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _wrapped_call_impl(self, *args, **kwargs)
   1516             return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517         else:
-> 1518             return self._call_impl(*args, **kwargs)
   1519 
   1520     def _call_impl(self, *args, **kwargs):

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py](https://localhost:8080/#) in _call_impl(self, *args, **kwargs)
   1525                 or _global_backward_pre_hooks or _global_backward_hooks
   1526                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527             return forward_call(*args, **kwargs)
   1528 
   1529         try:

[/usr/local/lib/python3.10/dist-packages/torch/nn/modules/loss.py](https://localhost:8080/#) in forward(self, input, target)
   1177 
   1178     def forward(self, input: Tensor, target: Tensor) -> Tensor:
-> 1179         return F.cross_entropy(input, target, weight=self.weight,
   1180                                ignore_index=self.ignore_index, reduction=self.reduction,
   1181                                label_smoothing=self.label_smoothing)

[/usr/local/lib/python3.10/dist-packages/torch/nn/functional.py](https://localhost:8080/#) in cross_entropy(input, target, weight, size_average, ignore_index, reduce, reduction, label_smoothing)
   3051     if size_average is not None or reduce is not None:
   3052         reduction = _Reduction.legacy_get_string(size_average, reduce)
-> 3053     return torch._C._nn.cross_entropy_loss(input, target, weight, _Reduction.get_enum(reduction), ignore_index, label_smoothing)
   3054 
   3055 

ValueError: Expected input batch_size (396) to match target batch_size (1170).
```


The problem was that I was accidentily sending the classifier labels to the model with a language modeling head. Not sure why this misformatted the batch size, but after removing there is no problem.

From this
`hidden = self.plm(input_ids=input_ids, attention_mask=attention_mask, labels=labels).hidden_states[-1]`
to
`hidden = self.plm(input_ids=input_ids, attention_mask=attention_mask).hidden_states[-1]`