Sharded DDP returns extra predictions EvalPrediction
Author: AADeLucia
Pull request: no

## Environment info
- `transformers` version: 4.17.0.dev0
- Platform: Linux-5.13.0-28-generic-x86_64-with-glibc2.17
- Python version: 3.8.11
- PyTorch version (GPU?): 1.10.2 (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?: Yes
- Using distributed or parallel set-up in script?: Yes, trying DataParallel and ShardedDDP

### Who can help
I'm using Trainer in a multi-GPU setup for training a custom model. Tagging @sgugger for Trainer.

## Information
I'm using a custom Multi-instance learning model with a BERTweet submodule. 

The problem arises when using:
* [N] the official example scripts: (give details below)
* [Y] my own modified scripts: (give details below)

The tasks I am working on is:
* [N] an official GLUE/SQUaD task: (give the name)
* [Y] my own task or dataset: (give details below)

## To reproduce

Settings passed to Trainer:
```bash
python -m torch.distributed.launch --nproc_per_node 2 "${MINERVA_HOME}/src/evaluation/mil_clean.py" \
  --run_name "finetune-shuffle" \
  --output_dir "${OUTPUT_DIR}" \
  --logging_dir "${LOG_DIR}" \
  --per_device_train_batch_size 1 \
  --per_device_eval_batch_size 30 \
  --learning_rate 1e-5 \
  --do_train \
  --do_eval \
  --do_predict \
  --logging_strategy "steps" \
  --log_on_each_node 0 \
  --logging_steps 50 \
  --log_level "info" \
  --save_strategy "steps" \
  --save_steps 1000 \
  --evaluation_strategy "steps" \
  --eval_steps 1000 \
  --optim "adamw_torch" \
  --load_best_model_at_end True \
  --metric_for_best_model "f1" \
  --num_train_epochs 10 \
  --dataloader_num_workers 2 \
  --dataloader_drop_last True \
  --num_tweets_per_day 100 \
  --shuffle_samples True \
  --seed 42 \
  --gradient_accumulation_steps 8 \
  --finetune_instance_model True \
  --dataloader_pin_memory False \
  --sharded_ddp "zero_dp_3"
```

```python
def compute_metrics(eval_prediction):
    probs, label_ids = eval_prediction  # Model returns probabilities instead of logits
    probs = probs.reshape(-1)
    predictions = (probs > 0.5).astype(np.uint8)
    try:
        prec, recall, f1, support = precision_recall_fscore_support(label_ids, predictions, zero_division=0, average="weighted")
    except ValueError as err:
        logger.error(f"{err=}\n{eval_prediction.predictions.shape=}\n{eval_prediction.label_ids.shape=}")
        exit(1)
    return {
        "precision": prec,
        "recall": recall,
        "f1": f1
    }
```

Output with ShardedDDP. No error is thrown with DataParallel so I think there is a collation issue. 
```bash
eval_prediction.predictions.shape=(14, 30)
eval_prediction.label_ids.shape=(378,)
```

Also, I saw "padding" mentioned in the Trainer documentation, but I printed out all the `predictions` values and did not see anything other than my probabilities.

> If your predictions or labels have different sequence length (for instance because you’re doing dynamic padding in a token classification task) the predictions will be padded (on the right) to allow for concatenation into one array. The padding index is -100.


## Expected behavior
I expect DataParallel and Distributed Data Parallel setups to send the same number of predictions to `compute_metrics`


Comments:
There is little we can do to help without knowing the code that is run and having a clear reproducer.

Ah sorry for the lack of detail, here is my code. 

https://gist.github.com/AADeLucia/86d5d75dde878cfa3dbd38a0230be2cf

Oh you're using custom models. Are they all returning loss then logits if labels are provided, logits only if no labels are provided? Also are the logits tensors with batch first dimension? While the `Trainer` works with other models, it has been optimized for Transformers models, so it has certain expectations.

It might also be specifically linked to sharded DDP. If you just used distributed training, is there still the issue?

I tried without the sharding and I received the same error (quick test with only 100 examples)

```bash
***** Running Evaluation *****
  Num examples = 100
  Batch size = 30

ERROR:root:err=ValueError('Found input variables with inconsistent numbers of samples: [100, 120]')4s/it]
eval_prediction.predictions.shape=(4, 30)
eval_prediction.label_ids.shape=(100,)
ERROR:root:err=ValueError('Found input variables with inconsistent numbers of samples: [100, 120]')
eval_prediction.predictions.shape=(4, 30)
eval_prediction.label_ids.shape=(100,)
```

So far I haven't omitted passing in labels, so it always assumes labels are there to compute loss. Logits are tensors with batch first dimension.  

Would it be best if I did something like this, using a huggingface output class?
```python
 from transformers.file_utils import ModelOutput
def forward(*****):
   output = {}
   output["logits"] = logits
   if labels are not None:
      output["loss"] = calculate loss
  return ModelOutput(output)
```


Not sure it will change anything to switch to a dictionary as long as you put your outputs in the right order (loss first if present) in your tuple. The `Trainer` can handle both. I'm curious to see what's being badly accumulated, could you try printing the shapes at the end of the forward pass of your model so we can try to figure out what's going on?

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.