Enable different torch dtype in sub models

[zucchini-nlp]

What does this PR do?

Fixes #33997. Enables users to use different torch dtypes for each of sub config. For ex load the vision model in full precision and the text model in half precision

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[zucchini-nlp]

if users passes one torch dtype as before, we just use it in all sub-configs. Otherwise a user can either set directly dtypes in configs before loading the model, or indicate a dict torch_dtype when loading similarly to attn_implementation_dispatch

[ArthurZucker]

Thanks, can you add a test showcasing an example usage of this (for example with Llava!)

[ArthurZucker]

This is interesting but I am not sure we have everything ready:

• does it work with the keep-in-float32 attribute as well?
• does it work well with model that have enforced param with dtypes?
  (Some vision models have this! ) So maybe a little bit of testing is missing
  Makes a lot of sense otherwise!

[zucchini-nlp]

Yeah, this one needs time and I'll come back after the model releases to make sure it works in all cases. Currently it has weird behavior in nested configs where a general text config has an attribute vision_config

[zucchini-nlp]

Conditions to make the dtype dispatch correctly: use _from_config to init the submodule inside a bigger model as ChameleonVQVAE._from_config. That way the default dtype is set before initilizing the weights on each submodule, otherwise they are not set

The current design support setting dtype via dict when loading or directly via changing the config. The dict should have the same keys as sub-configs and can have an empty string key which will be the main dtype. Main dtype is used for any module that is not part of sub-config, for ex mm_projector

Added more tests and verified it works when keep_in_fp32_modules is indicated. For the case when model's forcefully limit dtype of certain classes, I haven't found an actual model that does it so not sure how to test that

[ArthurZucker]

There seems to be a big breaking change no?

[zucchini-nlp]

Not breaking anymore, will need at least one approve to merge :)

[ArthurZucker]

Good to go but I'd rather wait for the next release 🤗

[zucchini-nlp]

cool, since release is done will rebase and merge