Seq2SeqTrainer.evaluate and predict don't yield the right number of predictions when num_return_sequences > 1
Author: antonioalegria
Pull request: no

### System Info

transformers: 4.31.0
accelerate: 0.21.0
python: 2.11.3
env: macOS 13.4.1

### Who can help?

@gante, I think, because this is related with generation

### 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

When calling evaluate or predict with `predict_with_generate` and `num_return_sequences` > 1, it does not pass the right amount of sequences to the `compute_metrics` function. It drops `num_return_sequences - 1` sequences in the last batch, in `Accelerator.gather_for_metrics`.

This does not happen when calling `model.generate`, which behaves as expected.

To reproduce run the following script:

```python
from transformers import (AutoModelForSeq2SeqLM, AutoTokenizer,
                          DataCollatorForSeq2Seq, GenerationConfig,
                          Seq2SeqTrainer, Seq2SeqTrainingArguments,
                          T5Tokenizer,BatchEncoding, PreTrainedTokenizer)
from transformers.utils import ModelOutput
from transformers.generation.utils import BeamSearchEncoderDecoderOutput
from datasets import Dataset, load_dataset

INPUT_COLUMN = "question"
TARGET_COLUMN = "answer"
MAX_INPUT_LENGTH = 256
MAX_TARGET_LENGTH = 256

dataset = load_dataset("gsm8k", "main", split="train[:38]")
model = AutoModelForSeq2SeqLM.from_pretrained("t5-small")
tokenizer=T5Tokenizer.from_pretrained("t5-small")
data_collator=DataCollatorForSeq2Seq(tokenizer, model=model, return_tensors="pt", padding="longest")
gen_config = GenerationConfig.from_pretrained("t5-small")
gen_config._from_model_config = None
gen_config.max_length = None
gen_config.min_length = None
gen_config.max_new_tokens = 256
gen_config.min_new_tokens = 1
gen_config.num_beams = 5

training_args=Seq2SeqTrainingArguments('.', predict_with_generate=True)

trainer = Seq2SeqTrainer(
    model=model,
    args=training_args,
    tokenizer=tokenizer,
    data_collator=data_collator,
    compute_metrics=lambda x: {"samples": x[0].shape[0]},
)

def prepare_data(examples: Dataset) -> BatchEncoding:        
    # Remove pairs where at least one record is none
    inputs = examples[INPUT_COLUMN]
    targets = examples[TARGET_COLUMN]
    
    model_inputs = tokenizer(inputs, max_length=MAX_INPUT_LENGTH, truncation=True)
    labels = tokenizer(text_target=targets, max_length=MAX_TARGET_LENGTH, truncation=True)
    model_inputs["labels"] = labels["input_ids"]
    
    return model_inputs

prepared_dataset = dataset.map(prepare_data, batched=True, remove_columns=[INPUT_COLUMN, TARGET_COLUMN])
dataset_len = len(prepared_dataset) # 38

gen_config.num_return_sequences = 1
metrics = trainer.evaluate(eval_dataset=prepared_dataset, num_beams = 5, generation_config=gen_config)
assert metrics["eval_samples"] == dataset_len

# THESE WILL FAIL -- THE NUMBER OF GENERATED SAMPLES WILL BE 70: 2*16 + 2*16 + 6 (last batch will discard the remaining 6 sequences)
gen_config.num_return_sequences = 2
metrics = trainer.evaluate(eval_dataset=prepared_dataset, num_beams = 5, generation_config=gen_config)
assert metrics["eval_samples"] == 2 * dataset_len # should be 76

# THESE WILL FAIL -- THE NUMBER OF GENERATED SAMPLES WILL BE 102: 3*16 + 3*16 + 6 (last batch will discard the remaining 32 sequences)
gen_config.num_return_sequences = 3
metrics = trainer.evaluate(eval_dataset=prepared_dataset, num_beams = 5, generation_config=gen_config)
assert metrics["eval_samples"] == 3 * dataset_len # should be 114
```

### Expected behavior

I would expect that the compute_metrics function would receive a tensor of shape (samples * num_return_sequences, max_len). Currently it receives a few less because the last batch gets half the sequences dropped in Accelerator.gather_for_metrics.

Comments:
It looks more like something in `accelerate`, so cc @muellerzr .

But @antonioalegria 

> . It drops num_return_sequences - 1 sequences in the last batch

Could you explain a bit more about this number? It doesn't seem corresponding to what you showed in the code snippet ..?

Apologies for not being clear.

Let's say you are generating from 100 input samples, `num_return_sequences` = 2 and eval batch size is 16.

You will have 6 full batches of 16, each generating 32 sequences, and a final batch of size 4. This final batch comes out of `model.generate` with 8 generated sequences but 4 of them are discarded in `Accelerator.gather_for_metrics`.

If you had `num_return_sequences` = 3, then the final batch would have originally 12 generated sequences, with 8 of them discarded in the end.

So final batch will always have the number of generated sequences equal to the last batch size.

Thanks! This should be solved via #27025