[Bugfix] Fix reloading of pixtral/llava configs
Author: kylesayrs
Pull request: yes

## Purpose ##
* Fix llava configs which attempt to override computed config values such as the `head_dim` on [mistral-community/pixtral-12b](https://huggingface.co/mistral-community/pixtral-12b)

## Related issues ##
* https://huggingface.co/mistral-community/pixtral-12b/discussions/24
* https://github.com/vllm-project/llm-compressor/issues/1127

## Background ##
Many issues arise when attempting to save and load Pixtral configs, namely that the pixtral model is saved with `text_config.head_dim=128`, but upon loading receives the value `text_config.head_dim=160`.

Through further inspection, it was found that this was due to an issue where the head_dim

MistralConfig
```
hidden_size=4096,
num_attention_heads=32,
head_dim = head_dim or hidden_size // num_attention_heads = 128
```

Pixtral 12B
```
hidden_size=5120,
head_dim=128,
num_attention_heads = None
```

When transformers configs are saved, a diff between the saving config and the default config is computed in order to reduce verbosity. When the Pixtral 12B config saves the `head_dim` attribute, it checks against the `MistralConfig` and finds that both values are the same (128), and therefore does not save the `head_dim` to disk.

However, when loading the same config using the `MistralConfig` class, the class recomputes `head_dim` using the saved `hidden_size=5120` (which differs from the default 4096), and returns the value `head_dim=160`. This incorrect calculation causes the model to be loaded incorrectly.

Note that for the pixtral model, `head_dim != hidden_size // num_attention_heads`, so this issue cannot be remedied by populating the `num_attention_heads` attribute of the pixtral config.

While I consider creating a dedicated `PixtralTextConfig` to be the most "correct" solution which reflects the differences in assumptions made by the pixtral and mistral models, implementing a config class in transformers which does not have a unique model associated with it requires additional work and testing. This solution is backwards compatible with existing pixtral configs and does not require intensive changes to configs.

## Changes ##
* Set `is_composition=True` for `LlavaConfig`, which forces all subconfig attributes to be saved to disk.
  * This means that `head_dim` is loaded and overrides the non-applicable calculation in MistralConfig\
* Add tests for reloading llava configs 

## Potential Follow-ups ##
* Because the pixtral model does not follow the assumptions made by the MistralConfig, the pixtral model should perhaps be given its own config.
  * For pixtral, `head_dim != hidden_size // num_attention_heads`
  * #36140
* Add `is_composition` flag to all configs with subconfigs
  * #36137

## Testing ##
```python3
from transformers import LlavaConfig

config = LlavaConfig.from_pretrained("mistral-community/pixtral-12b")
assert config.text_config.head_dim == 128

config.save_pretrained("tmp")
config = LlavaConfig.from_pretrained("tmp")
assert config.text_config.head_dim == 128  # previously, this would be overwritten by MistralConfig to be 160
```

Comments:
Hi @zucchini-nlp!

> I think we should find why the saved/loaded config inferred incorrect LM-config type and fix it the way, so that we can infer correct lm-config

This would likely involve having to modify the `PretrainedConfig` base class to initialize using the `LlavaConfig` scheme. Specifically,
https://github.com/huggingface/transformers/blob/b673c16cad81c71f70903a9a63f5b5f06014aa9e/src/transformers/configuration_utils.py#L820

Calls the init with no arguments, which causes `LlavaConfig` to default to a Llama text config
https://github.com/huggingface/transformers/blob/014047e1c8784c00e2a04cb04ffcecdd5cb23c16/src/transformers/models/llava/configuration_llava.py#L128-L132

This is an inevitable problem because `LlavaConfig` must have **some default** if no arguments are provided. Changing this default to `MistralConfig` will break all configs which default to `LlamaConfig`.

Your comment suggests modifying `PretrainedConfig.to_diff_dict` to specialize with init args. Maybe something like
```python3
init_args = {key: getattr(value, "model_type", None) for key, value in self.sub_configs}
class_config_dict = self.__class__(**init_args).to_dict() if not self.is_composition else {}
```

The only problem with this implementation is that this forces all configs to follow the particular scheme of using the `sub_configs` attribute and specializing using the `model_type` subattribute and nothing else. This would mean having to fix configs which do not follow this scheme like [`HybridClipConfig`](https://github.com/huggingface/transformers/blob/014047e1c8784c00e2a04cb04ffcecdd5cb23c16/examples/research_projects/jax-projects/hybrid_clip/configuration_hybrid_clip.py#L55) and likely a large suite of tests across all configs.

> Hey, thanks for finding the root cause. However I don't think we can solve this by just setting the flag. This flag is to be used for a bit different purposes and the naming a misleading

@zucchini-nlp Could you provide a little more context as to how this flag is used? Grepping the transformers codebase reveals that this attribute is only consumed by `PretrainedConfig.to_diff_dict` and is set to True for many nested configs similar to `LlavaConfig`.
https://github.com/search?q=repo%3Ahuggingface%2Ftransformers%20is_composition&type=code

There may be some upstream use that I am unaware of.

Thanks, now I see why the model types are not matching. The flag was introduces at first for models like Musicgen, when the sub-configs had no default values set and had to be indicated by users explicitly (in https://github.com/huggingface/transformers/pull/25237). 
https://github.com/huggingface/transformers/blob/014047e1c8784c00e2a04cb04ffcecdd5cb23c16/src/transformers/models/musicgen/configuration_musicgen.py#L200-L203


For Llava we didn't set it, because we have default values. But now since the same config class is used with various LM backbones, the defaults don't make much sense. I agree with the proposed fix and imo we'll need the flag to be added in all confgis that accept any `model_type` for sub-config. 

Can you propagate it to all VLMs like Llava and add a small test?

WIP comment: The cause of this issue is more subtle than I initially thought. There already exist mechanisms for catching diffs of non-default subconfigs that I initially missed, but they seem to not work well with computed config values.

https://github.com/huggingface/transformers/blob/6b550462139655d488d4c663086a63e98713c6b9/src/transformers/configuration_utils.py#L832

https://github.com/huggingface/transformers/blob/6b550462139655d488d4c663086a63e98713c6b9/src/transformers/configuration_utils.py#L1181

Mistral has default values
```
hidden_size=4096,
num_attention_heads=32,
self.head_dim = head_dim or hidden_size // num_attention_heads = 128
```

Pixtral has default values
```
hidden_size=5120,
head_dim=128,
num_attention_heads = None
```

Because the default head dim values match, the diff does not register and the head_dim is not written. But upon reloading, head_dim is recomputed using the new hidden_size, which results in the 160 value.

This maybe fixable by simply adding `num_attention_heads=40` to the pixtral config


@kylesayrs yeah, that's what I thought at first but the Pixtral model checkpoint actually has `num_attention_heads=32`. So it wouldn't work in that case. The reason is that for Pixtral `head_dim * num_attention_heads != hidden_size`

@zucchini-nlp I've updated with a WIP of what a pixtral text config might look like. Let me know what you think. If this approach isn't viable, we can fall back to `is_composition`

@zucchini-nlp I think the `is_composition` solution works, but it creates very verbose configs which may be unnecessary. To me, the problem is that the MistralConfig does not apply to the pixtral model, as the Mistral assumption that `head_dim = hidden_size // num_attention_head` is broken.

I still think that using `is_composition` is a good solution for now, as it ensures backwards compatibility with existing pixtral configs and doesn't involve creating a new "pixtral_text" config which does not correspond to a unique model (which afaict requires extra work to support with `transformers`)

@zucchini-nlp I've updated the PR with a clearer explanation of the issue and links to potential follow ups.

cc @ArthurZucker if you want to give a look

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

Will merge it, since it's not core and nice to have in the next release