Inconsistent shape of logits in `GenerateBeamDecoderOnlyOutput`
Author: kurzdev
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.48.3
- Platform: Linux-4.18.0-513.5.1.el8_9.x86_64-x86_64-with-glibc2.35
- Python version: 3.11.11
- Huggingface_hub version: 0.30.2
- Safetensors version: 0.5.3
- Accelerate version: 1.6.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.6.0+cu124 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: yes, across 4 GPUs with `accelerate launch --multi-gpu`
- Using GPU in script?: yes
- GPU type: NVIDIA A100-SXM4-80GB

### Who can help?

@zucchini-nlp @stevhliu

### 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

```py
import torch
from torch.utils.data import DataLoader
from transformers import AutoProcessor, LlavaNextForConditionalGeneration


BATCH_SIZE = 4
NUM_BEAMS = 5

model = LlavaNextForConditionalGeneration.from_pretrained(
    "llava-hf/llama3-llava-next-8b-hf",
    torch_dtype=torch.bfloat16,
    low_cpu_mem_usage=True,
    attn_implementation="flash_attention_2",
)
processor = AutoProcessor.from_pretrained(
    "llava-hf/llama3-llava-next-8b-hf",
    model_max_length=4096,  # https://llava-vl.github.io/blog/2024-04-30-llava-next-video/
)

dataset = ...  # Some custom VQAv2 dataset
loader = DataLoader(
    dataset,
    batch_size=BATCH_SIZE,
    num_workers=2,
    pin_memory=True,
    shuffle=False,
)

for batch in loader:
    features = processor(
        images=batch["image"],
        text=batch["prompt"],
        return_tensors="pt",
        padding="max_length",
        truncation=True
    )

    outputs = model.generate(
        max_new_tokens=128,
        num_beams=NUM_BEAMS,
        length_penalty=0,
        input_ids=features["input_ids"],
        attention_mask=features["attention_mask"],
        pixel_values=features["pixel_values"],
        image_sizes=features["image_sizes"]
        pad_token_id=processor.tokenizer.eos_token_id,
        return_dict_in_generate=True,
        output_logits=True
    )

    stacked_logits = torch.stack(outputs.logits)
    print(stacked_logits.shape)  # Should be (seq_length, batch_size, vocab_size), is (seq_length, batch_size * num_beams, vocab_size)
```

### Expected behavior

I am using `llava-hf/llama3-llava-next-8b-hf` for generation on [VQAv2](https://visualqa.org/download.html). When generating in batches, the actual shape of the logits is `(seq_length, batch_size * num_beams, vocab_size)` instead of `(seq_length, batch_size, vocab_size)` as is stated [in the documentation](https://huggingface.co/docs/transformers/en/internal/generation_utils#transformers.generation.GenerateBeamDecoderOnlyOutput).
So for greedy decoding (which is equivalent to `num_beams=1`) and a batch size of 4, I get logits of shape `(n, 4, m)`, whereas for `num_beams=5`, I get `(n, 20, m)`.

I think this is because for each beam and element in the batch, a separate tensor is added to `raw_logits` in `_beam_search` within the `GenerationMixin`:
https://github.com/huggingface/transformers/blob/v4.51.3/src/transformers/generation/utils.py#L3915-L3928

I'm not sure if this is a bug in the implementation (which I'd assume) or an error in the documentation, hence the double-ping.
Either way, my goal is to obtain the logits for the actually generated sequence out of all beams. So if this is just an inconsistency in the documentation, I'd really appreciate it if you could point me in the direction of how to use the resulting shape to obtain the logits for the actually generated sequence per batch!

Comments:
@kurzdev hey! yes, it is expected because we return all scored for each beam in generation. You can take a look at this example of how to obtain scores of the selected beam -> https://discuss.huggingface.co/t/announcement-generation-get-probabilities-for-generated-output/30075

@zucchini-nlp hey, thanks for the swift response and the reference - will look into it! that does make sense! 😊
in that case, I think we should update the documentation to make this clearer in the future.

cc @gante and @manueldeprada , I am not sure if we have it documented anywhere

Yup the docs are incorrect: it should return the logits for all beams

#37929 fixes it :)