Correct llava mask & fix missing setter for `vocab_size`
Author: fxmarty
Pull request: yes

This PR fixes llava mask in the generation case.

`torch.cat((attention_mask, extended_mask), dim=-1)` was very wrong and it is a miracle we were getting meaningful generation for batch generation. The issue is that this disregards the attention_mask used during the first forward pass, that is custom. Wrong past key values are used.

It is also not clear to me why this custom masking is handled in the forward and not in GenerationMixin cc @gante (we are doing unnecessary operations at each forward to retrieve the correct attention_mask, it could just be an output of the model & updated in GenerationMixin)

Fixes https://github.com/huggingface/transformers/issues/28184

Comments:
The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_29389). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

> It is also not clear to me why this custom masking is handled in the forward and not in GenerationMixin cc @gante (we are doing unnecessary operations at each forward to retrieve the correct attention_mask, it could just be an output of the model & updated in GenerationMixin)

@fxmarty We do that by default (in [this function](https://github.com/huggingface/transformers/blob/9c5e560924c4868033cc9ac650bee80dc4049158/src/transformers/generation/utils.py#L647)) 🤔 

@zucchini-nlp since you're interested in multimodal models: after this PR gets merged, can you inspect why llava (and related models) need this custom attention mask handling at generation time?

@gante , I've been affected by this also while trying to make speculative decoding work with VLMs. I'll briefly outline what I found.

The custom mask is needed in Llava because it concatenated image embeds and text, which is further used as `inputs_embeds` for the LLM part. Thus after pre-fill they require attention mask of much larger length than what we got initially. When we update `model_kwargs` in `generate`, we assume one token is added to prev attention mask. This is not the case for Llava, since the attention mask holds values for textual part only, so the `forward` adds masks for the image part every call.

Comparing to other models, I have counted only 3 soft prompt VLMs in `transformers` (Kosmos-2, Blip, Llava). The other two models do not have this issue, because they have custom `generate` where embeds are concatenated, new attention masks are prepared manually and only then `text_model.generate()` is called for further handling. But Llava calls `composite_model.generate()` directly.

We can either fix Llava to call `text_model.generate()` for consistency with others or update `model_kwargs` depending on `past_key_values` length in `generate`. I was also thinking that if the second case, we can get rid of using `_prepare_token_type_ids` and `_prepare_attention_mask` in speculative decoding. 

It's up to you to decide if we need any changes, I have no idea how this can affect library-wide 😄 

> We can either fix Llava to call text_model.generate() for consistency with others or update model_kwargs depending on past_key_values length in generate. I was also thinking that if the second case, we can get rid of using _prepare_token_type_ids and _prepare_attention_mask in speculative decoding.

@zucchini-nlp I'd say to go with "call text_model.generate() for consistency", if that is feasible and results in a clean interface. `generate` is still quite rigid at the moment, so handling more cases will slow down ongoing projects. After we refactor `generate` to make it more flexible, we can revisit this decision :)

Llava slow tests all pass (apart from `tests/models/llava/test_modeling_llava.py::LlavaForConditionalGenerationModelTest::test_cpu_offload` which does not pass on main either).

I added a setter to fix the regression in https://github.com/huggingface/transformers/pull/29586#discussion_r1535058927 (`RUN_SLOW=1 CUDA_VISIBLE_DEVICES=1 pytest tests/models/llava -s -vvvvv -k "test_small_model_integration_test"` failing on main)

I wonder if llava is working on the release version? maybe not

@amyeroberts We need to set `final_embedding` to zero for https://github.com/huggingface/transformers/blob/aa17cf986f0761382f3d6e591e985a42671c3fb7/src/transformers/models/llava/modeling_llava.py#L449

to do its job! Up to now it did not do what was expected.

What I meant in

> It is also not clear to me why this custom masking is handled in the forward and not in GenerationMixin cc @gante (we are doing unnecessary operations at each forward to retrieve the correct attention_mask, it could just be an output of the model & updated in GenerationMixin)

is that this (i.e. the attention_mask reconstruction in the decode) could be avoided altogether if the mask was rather handled in generate, which is currently not the case in the implementation

edit: slow tests for llava pass