Attn implementation for composite models
Author: zucchini-nlp
Pull request: yes
Labels: name: run-slow

# What does this PR do?

Related to #32221 and fixes #30565. As described in the issue, currently MultiModal models set attn implementation only in the LM backbone, and not in vision backbone. While working on it, I found another bug. Specifically, using a property to set SDPA flag to VLM class doesn't work because we don't know what is `self.LM` before init the VLM. I tried class property, but that would throw error that `self has no attribute LM`. 

This PR makes a few modifications to how attn impl is set in modeling to fix the above issues.

More precisely:

- If the model is composite, i.e. has 2 or more PreTrainedConfigs as part of the main config (text/vision config), then we set the requested attention implementation to each of the sub-configs as a dict. For ex, in LLaVa case:
```python
model.from_pretrained(id, attn_implementation="sdpa")

# internally we do the following, where the second line dispatches SDPA on LM backbone
# if we can't dispatch, an error is raised as it should be when loading an LLM with unsupported attention (i.e. llama doesn't support SDPA)
self.config._attn_implementation = {"text_config": "sdpa", "vision_config": "sdpa"}
AutoModelForCausalLM._from_config(config.text_config, attn_implementation=config._attn_implementation["text_config"])
```

- If user just loads the composite model without any explicit attention, then we dispatch SDPA in any of the sub-models whenever possible. In case any/none of the vision/text models support SDPA, we silently fallback to eager. No warnings, no errors, which is in line with non-composite models
```python
# suppose LM backbone doesn't support SDPA here, then we'll use eager silently
# but dispatch SDPA on vision backbone
model.from_pretrained(id, attn_implementation=None)
```

- The users are also free now to set different attentions in different backbones. We can set eager on vision model and sdpa on LM. That can be done if one passes a dict `attn_implementation`. And the dict keys have to be identical to the ModelConfig attribute names for sub-configs. In other words, if the config has config.text_config, then the dict key must be "text_config" to correctly dispatch

```python
model.from_pretrained(id,  {"text_config": "sdpa", "vision_config": "eager"})
```
- For non-composite models nothing changes and everything works as before


NOTE:
I had a few rare cases where the model was not really composite but had several sub-configs. For example DBRX, and I decided to get rid of the sub-configs there. I am not sure what is the right way to deprecate that, so any comments on that welcome
Also, some models have nested configs, where a vision config is part of text config (Qwen2-VL). Only three models are like that, mostly because I didn;t review properly when the model was added and didn't insist on changing stuff. For those models, we don't consider them as composite and just leave it working as it was


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

Regarding

```
MODEL_MAPPING.get(type(config.text_config), None)
            if text_model_cls is not None:
```
it could probably reuse

```
def _get_model_class(config, model_mapping):
```

But both will be somehow slow as you mentioned, and sometimes give unexpected errors (because loading a modeling module file). But we might have someway to avoid that.

@ydshieh @qubvel I made some changes after our last discussion that changing class attr is not good. So now we don't change the attr per se, but we change the config's `attn_impl` manually while setting sdpa. I checked that it works this way also. The last question is whether to set it to SDPA when both sub-configs support it or at least one supports it. 

I will add more tests if we agree on the design


And also, regarding
> But both will be somehow slow as you mentioned, and sometimes give unexpected errors (because loading a modeling module file). But we might have someway to avoid that.

unfortunately we can't do that, cause it raises error if key is not in the dict and we have not auto-mappable models. So we rather try to get and if not simple continue. 


> unfortunately we can't do that, cause it raises error if key is not in the dict and we have not auto-mappable models. So we rather try to get and if not simple continue.

but why not try except around the usage of `_get_model_class`? I don't feel very strong here, but just it is something well designed in order to get the model class which is good if we can reuse it.

Hmm, not sure how good if to use a `try except`, can it be another source of hidden bugs even if we use `except AttributeError`? I'm okay with any way, as long as it's consistent with the library and doesn't cause bugs

As we discussed internally, I made the loading logic model-agnostic. So now we check if the config contains any `PretrainedConfig` inside and set attn implementation to each `PretrainedConfig`. This solved vision-encoder-decoder problem we had in slack and makes us more versatile for possible future multimodal models.

I added some tests in vision enc-dec models where only one sub-config supports SDPA and when both support. Also added tests in Idefics (which were skipped earlier) to verify that LM has its SDPA dispatched even though other modules do not.

TODO in the next few hours or tomorrow:
- Add similar tests to encoder-decoder models, not only vision enc-dec
- Accept from users a config with separate `attn_implementation` already set, such that `attn_implementation={"vision_config": "sdpa", "text_config": "eager"}`

I made final changes to the PR and made sure all tests are passing and all composite models were modified. The main points are:

- we have a `Is-composite` flag for models so that we know if we want to dispatch on each sub-config or not. The flag was needed for models like Chameleon that have a sub-config, but it's for VAE which will not supports SDPA i guess
- attn implementation can be set as prev ("sdpa") or as a dict in case of composite models ({"encoder": "sdpa", "decoder": "eager"}). For composite models, in first case we try to dispatch requested attn in all sub-models and raise error if not supported. In second case we dispatch the different attn for each sub-model as requested and also raise error if not supported
- Idefics and MusicGen models can get rid of custom code for attn implementation now
- For testing I had to add another flag `is_multimodal` (maybe will change to `is_composite`) and check each sub-config's attn implementation is set correctly. Also added tests in enc-dec models
- For some models which were not automappable, I had to add them to mapping so that we can get from config to class, and obtain `supports_sdpa` flag


One last thing to consider is that this might result in a bit of slow-down as we try to do automapping for each sub-model and dispatch separately. I have another idea, we can stop dispatching in each sub-model and instead simply replace "str" attn implementation with "dict" as `attn_implementation = {"encoder": requested_attn, "decoder": requested_attn}`. Then inside modeling we could do `AutoModel.from_config(config, attn_impl = config._attn_impl["encoder"])`. I testes with a few models and verified it works same way as current PR but did not introduce those changes yet. First, I want to get a general feedback if the PR is doing what we want. cc @amyeroberts, would like to know if I addressed your concerns about hidden magic code

Made updates and removed automapping, seems we don't need it to dispatch in submodules. Added more tests, specifically for FA2 and for the new way of setting attn impl (as a dict in composite models). All the tests I added/touched are passing for me locally + non-slow tests in all models.

Some FA2/SDPA tests are failing but those seem to be failing disregarding the current changes. Requesting re-review, failing tests are not related and should have been skipped

UPDATE: Ran slow tests on LLaVA, Paligemma, Llama, Idefics1 and 2. All is green

Will [this warning](https://huggingface.co/Qwen/Qwen2-VL-7B-Instruct/discussions/23) stop appearing after this PR is merged?

@ncoop57 not really in the current state, because we always call `from_config(attn_implementation="my_attn_impl")` to initialize a sub-component and we don't indicate `dtype`. Actually, now that I look into that, seems like we could call simply `SubModule(config)` instead of `_from_config`. The only thing is that directly calling class won't work with deepspeed, according to code. I'll check this out, should be also fixed in scope of this PR

TODO: New VLMs were added while PR was in progress, so we need to propagate new design to:
- [ ] Qwen2-VL
- [ ] LLava-Onevision
- [ ] Idefics2
- [ ] Some more coming soon


In the meanwhile, @amyeroberts or @ArthurZucker , can I get a review on general design ?

>  i.e. if it's compatible in one block but not the other an error will still be raised?

Yes this is the behavior now. Sorry I forgot to update PR description so let me describe the expected behavior.

For non-composite models nothing changes and the attention is dispatched same way as before. For composite models:

- if user just loads the model without any explicit attention (`Model.from_pretrained(model_id)`), then we dispatch SDPA in any of the sub-models whenever possible. In case any/none of the vision/text models support SDPA, we silently fallback to eager. No warnings, no errors, which is in line with non-composite models
- if user specifies required attention, we check if all of the sub-models support it. An error will be raised on sub-model level if the attention is not yet supported, so instead of `Llava doesn't support FA2`, we'll raise `CLIP doesn't support FA2`
- if user specifies attention for each sub-model when loading, as a dict (`Model.from_pretrained(model_id, attn_implementation={"vision_config": "eager", "text_config": "sdpa"})`), then we dispatch required attention to indicated models. We raise an error if any of the requested attentions is not supported by the indicated sub-model. When using attn implementation as a dict the user has to indicate attn for all sub-models, no defaults will be available. And the dict keys have to be identical to the `ModelConfig` attributes for sub-configs. In other words, if the config has `config.text_config`, then the dict key must be `text_config` to correctly dispatch

> * if user specifies required attention, we check if all of the sub-models support it. An error will be raised on sub-model level if the attention is not yet supported, so instead of `Llava doesn't support FA2`, we'll raise `CLIP doesn't support FA2`

@zucchini-nlp Great, thanks for explaining. Behaviour all sounds good to me! 

I think now this is ready for review, @ArthurZucker 

I rebased main and updated all new models, addressed some comments from Amy and added a few more tests

Related to #33953 as well, will have a look, but having to change modeling code means we are doing something wrong as it should be a layer of abstraction! 

Cool thanks! The modeling is mostly changed to get the attn implementation from dict and for IDEFICS (1 and 2), which fall out of standards and enforce hard requirements on sdpa. Hope we can get this merged soon, as it is getting more of a problem when new quirky models are added

@ArthurZucker I made changes as you suggested but I want to warn that it still needed a few extra logic. Making this PR revealed soo many issues. So we have two options now to indicate attn for composite models:

**Attn implementation as a dict:**
The prev solution when we simply copy paste the user-requested attn in all sub-configs, and form a dict. So that the general `config.attn = Dict(text_config: None, vision_config: None)`
+ no need for flag whether attn implementation is set or not, because when we have a dict we always use the very first requested attn in all sub-configs. We just need to save it as a dict once and then go back to that dict to "get()"
+ we don't need to unset attn implementation when saving so one can set a different attn when loading any model back.
- BUT we change a lot of model code to pass in explictly (attn_implementation=config._attn_implementation["text_config"]) and we need to account for cases like Chameleon or DBRX where the sub-configs don't even need attn

**Attn implementation directly propagated in corresponding configs:**
Current solution when we simply propagate the attn to all configs we can find, and don't touch the general config
+ no need to change any modeling code, except for general clean up
- BUT we need a flag attn_implementation_autoset  so we don't dispatch attn on the same config twice. Also we would need to unset the flag when saving. See below for reason

Imagine we have a composite or nested config. Currently we set attn in three places: `from_pretrained`, `from_config` and `init`. When we don't require any attn (`None`) and load with `model.from_pretrained(id)`, internally we will call `autoset_attn`  and set SDPA to the general config and the sub-configs will get `None` propagated. But then the same method `autoset_attn`  is called second time when we init the class, and this time we'll try to propagate SDPA to all sub-configs (because general config has SDPA already set). If any sub-module has no SDPA support, an error will be raised as it was requested directly. From user perspective they did nothing wrong, but they see an error.
So the decision is to either stop autosetting attn for each config once it was set, which is happening now. Or in prev version I had attn implementations saved as dict, which basically did the same. When we got a dict attn, we had nothing to set and simply returned

If we find a clean way to auto_set attn implementation would be better in the futur! 

merging tomorrow, the tests pass locally in most cases except fro the cases not touched by this PR. For ex some FA2/sdpa equivalence tests are failing for me in `main` branch