The length used by length_penalty during beam_search is not correct when input batch size > 1

### System Info Python 3.9.2 transformers 4.30.2 ### Who can help? @gante ### Information - [X] The official example scripts - [X] My own modified scripts ### Tasks - [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...) - [x] My own task or dataset (give details below) ### Reproduction ``` from transformers import BloomTokenizerFast, BloomForCausalLM import torch tokenizer = BloomTokenizerFast.from_pretrained('bigscience/bloomz-7b1') tokenizer.padding_side = "left" model = BloomForCausalLM.from_pretrained('bigscience/bloomz-7b1') model = model.cuda().half().eval() def predict(model, text): token_out = tokenizer(text, padding=True, return_tensors="pt") input_ids = token_out.input_ids.cuda() attention_mask = token_out.attention_mask.cuda() generate_kwargs = dict(max_new_tokens=128, do_sample=False, num_beams=2, length_penalty=3.0) ori_out = model.generate(input_ids=input_ids, attention_mask=attention_mask, **generate_kwargs) print('ori input shape: ', input_ids.shape) print('ori output shape: ', ori_out.shape) return tokenizer.batch_decode(ori_out, skip_special_tokens=True) text = ["Explain backpropagation in neural networks."] baseline = predict(model, text)[0] text = ["Explain backpropagation in neural networks.", "What is 1 + 1?"] good_case = predict(model, text)[0] text = ["Explain backpropagation in neural networks.", "The Basilica of the Sacred heart at Notre Dame is beside to which structure? The Basilica of the Sacred heart at Notre Dame is beside to which structure? The Basilica of the Sacred heart at Notre Dame is beside to which structure?"] bad_case = predict(model, text)[0] print(f'baseline: \n{baseline}\ngood_case: \n{good_case}\nbad_case: \n{bad_case}\n') ``` The script results: ``` baseline: Explain backpropagation in neural networks. backpropagation is the process by which the error is propagated backwards through the network from the output layer to the input layer, and then back to the output layer aga…